update StarHead skill (#278)

TimoKorinth · web-flow · commit 18518fea9583 · 2025-01-29T10:24:55.000+01:00
* add shop info

* update prompt

* try another approach

* update function naming

* fix naming

* first try to integrate vision skill in starhead

* update prompts and image generation skill

* refactor after rebase

* remove vision functions for now

* update template
diff --git a/skills/image_generation/main.py b/skills/image_generation/main.py
@@ -39,36 +39,40 @@ async def execute_tool(
 
         if tool_name == "generate_image":
             prompt = parameters["prompt"]
-            if self.settings.debug_mode:
-                await self.printr.print_async(f"Generate image with prompt: {prompt}.", color=LogType.INFO)
-            image = await self.wingman.generate_image(prompt)
-            await self.printr.print_async(
-                        "",
-                        color=LogType.INFO,
-                        source=LogSource.WINGMAN,
-                        source_name=self.wingman.name,
-                        skill_name=self.name,
-                        additional_data={"image_url": image},
-                    )
-            if image:
-                function_response = "Here is an image based on your prompt."
-
-                if self.save_images:
-                    image_path = path.join(
-                        self.image_path,
-                        f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}_{prompt[:40]}.png"
-                    )
-                    image_response = requests.get(image)
-
-                    if image_response.status_code == 200:
-                        with open(image_path, 'wb') as file:
-                            file.write(image_response.content)
-
-                        function_response += f" The image has also been stored to {image_path}."
-                        if self.settings.debug_mode:
-                            await self.printr.print_async(f"Image displayed and saved at {image_path}.", color=LogType.INFO)
-
+            await self.generate_image(prompt)
+            function_response = "Here is an image based on your prompt."
         return function_response, instant_response
+    
+    async def generate_image(self, prompt: str) -> str:
+        if self.settings.debug_mode:
+            await self.printr.print_async(f"Generate image with prompt: {prompt}.", color=LogType.INFO)
+
+        image = await self.wingman.generate_image(prompt)
+        await self.printr.print_async(
+                    "",
+                    color=LogType.INFO,
+                    source=LogSource.WINGMAN,
+                    source_name=self.wingman.name,
+                    skill_name=self.name,
+                    additional_data={"image_url": image},
+                )
+        if image:
+            function_response = "Here is an image based on your prompt."
+
+            if self.save_images:
+                image_path = path.join(
+                    self.image_path,
+                    f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}_{prompt[:40]}.png"
+                )
+                image_response = requests.get(image)
+
+                if image_response.status_code == 200:
+                    with open(image_path, 'wb') as file:
+                        file.write(image_response.content)
+
+                    function_response += f" The image has also been stored to {image_path}."
+                    if self.settings.debug_mode:
+                        await self.printr.print_async(f"Image displayed and saved at {image_path}.", color=LogType.INFO)
 
     async def is_waiting_response_needed(self, tool_name: str) -> bool:
         return True
diff --git a/skills/star_head/default_config.yaml b/skills/star_head/default_config.yaml
@@ -31,6 +31,9 @@ prompt: |
   Your job is to find good trading routes for the user based on his/her ship, current location and available budget.
   The user can also ask you about details of specific ships, components, weapons, and more.
   You always use the tools available to you to retrieve the required information and to provide the user with the information.
+  Do not try to determine a parent object by yourself, always ask the user about it.
+  Don't provide possible values of a parameter, always ask the user about it.
+  The currency is always 'aUEC' spelled 'Alpha UEC'.
 custom_properties:
   - hint: The URL of the StarHead API.
     id: starhead_api_url
diff --git a/skills/star_head/main.py b/skills/star_head/main.py
@@ -37,6 +37,9 @@ def __init__(
         self.celestial_objects = []
         self.celestial_object_names = []
         self.quantum_drives = []
+        self.shops = []
+        self.shop_names = []
+        self.shop_parent_names = []
 
     async def validate(self) -> list[WingmanInitializationError]:
         errors = await super().validate()
@@ -78,6 +81,19 @@ async def _prepare_data(self):
             "vehiclecomponent", {"typeFilter": 8}
         )
 
+        self.shops = await self._fetch_data("shop")
+        self.shop_names = [shop["name"] for shop in self.shops]
+
+        # Remove duplicate shop names
+        self.shop_names = list(dict.fromkeys(self.shop_names))
+
+        self.shop_parent_names = [
+            shop["parent"]["name"] for shop in self.shops if shop["parent"]
+        ]
+
+        # Remove duplicate parent names
+        self.shop_parent_names = list(dict.fromkeys(self.shop_parent_names))
+
     async def _fetch_data(
         self, endpoint: str, params: Optional[dict[str, any]] = None
     ) -> list[dict[str, any]]:
@@ -113,6 +129,17 @@ async def execute_tool(
             function_response = await self._get_best_trading_route(**parameters)
         if tool_name == "get_ship_information":
             function_response = await self._get_ship_information(**parameters)
+        if tool_name == "get_trading_information_of_specific_shop":
+            function_response = await self._get_trading_information_of_specific_shop(
+                **parameters
+            )
+        if tool_name == "get_trading_shop_information_for_celestial_objects":
+            function_response = (
+                await self._get_trading_shop_information_for_celestial_objects(
+                    **parameters
+                )
+            )
+
         return function_response, instant_response
 
     async def is_waiting_response_needed(self, tool_name: str) -> bool:
@@ -135,9 +162,9 @@ def get_tools(self) -> list[tuple[str, dict]]:
                                     "type": "string",
                                     "enum": self.celestial_object_names,
                                 },
-                                "moneyToSpend": {"type": "number"},
+                                "money_to_spend": {"type": "number"},
                             },
-                            "required": ["ship", "position", "moneyToSpend"],
+                            "required": ["ship", "position", "money_to_spend"],
                         },
                     },
                 },
@@ -159,10 +186,94 @@ def get_tools(self) -> list[tuple[str, dict]]:
                     },
                 },
             ),
+            (
+                "get_trading_information_of_specific_shop",
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_trading_information_of_specific_shop",
+                        "description": "Gives trading information about the given shop, like which commodities you can sell or buy and for which price. If the name of the shop is not unique, you have to specify the celestial object.",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "shop": {"type": "string", "enum": self.shop_names},
+                            },
+                            "required": ["shop"],
+                        },
+                    },
+                },
+            ),
+            (
+                "get_trading_shop_information_for_celestial_objects",
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_trading_shop_information_for_celestial_objects",
+                        "description": "Gives trading information about the given celestial object, like which commodities you can sell or buy at which shop and for which price. All shops with shop items on that celestial object will be returned.",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "celestial_object": {
+                                    "type": "string",
+                                    "enum": self.celestial_object_names,
+                                },
+                            },
+                            "required": ["celestial_object"],
+                        },
+                    },
+                },
+            ),
         ]
 
         return tools
 
+    async def _get_trading_shop_information_for_celestial_objects(
+        self, celestial_object: str
+    ) -> str:
+        object_id = self._get_celestial_object_id(celestial_object)
+
+        if not object_id:
+            return f"Could not find celestial object '{celestial_object}' in the StarHead database."
+
+        shops = await self._fetch_data(f"shop?celestialObjectFilter={object_id}")
+
+        shop_items = {}
+        for shop in shops:
+            items = await self._fetch_data(f"shop/{shop['id']}/items")
+            for item in items:
+                item["pricePerItem"] = item["pricePerItem"] * 100
+                item["tradeType"] = (
+                    "Sold by store" if item["tradeType"] == "buy" else "The shop buys"
+                )
+            shop_items[f"{shop['parent']['name']} - {shop['name']}"] = items
+
+        shop_details = json.dumps(shop_items)
+        return shop_details
+
+    async def _get_trading_information_of_specific_shop(self, shop: str = None) -> str:
+        # Get all shops with the given name
+        shops = [
+            shop for shop in self.shops if shop["name"].lower() == shop["name"].lower()
+        ]
+
+        # Check if there are multiple shops with the same name
+        if len(shops) > 1:
+            return f"Multiple shops with the name '{shop}' found. Please specify the celestial object."
+
+        if not shops:
+            return f"Could not find shop '{shop}' in the StarHead database."
+
+        shop_items = {}
+
+        for shop in shops:
+            items = await self._fetch_data(f"shop/{shop['id']}/items")
+            for item in items:
+                item["pricePerItem"] = item["pricePerItem"] * 100
+            shop_items[shop["name"]] = items
+
+        shop_details = json.dumps(items)
+        return shop_details
+
     async def _get_ship_information(self, ship: str) -> str:
         try:
             response = requests.get(
@@ -177,7 +288,7 @@ async def _get_ship_information(self, ship: str) -> str:
         return ship_details
 
     async def _get_best_trading_route(
-        self, ship: str, position: str, moneyToSpend: float
+        self, ship: str, position: str, money_to_spend: float
     ) -> str:
         """Calculates the best trading route for the specified ship and position.
         Note that the function arguments have to match the funtion_args from OpenAI, hence the camelCase!
@@ -195,7 +306,7 @@ async def _get_best_trading_route(
             "startCelestialObjectId": celestial_object_id,
             "quantumDriveId": qd["id"] if qd else None,
             "maxAvailablScu": cargo,
-            "maxAvailableMoney": moneyToSpend,
+            "maxAvailableMoney": money_to_spend,
             "useOnlyWeaponFreeZones": False,
             "onlySingleSections": True,
         }
@@ -215,7 +326,7 @@ async def _get_best_trading_route(
         if parsed_response:
             section = parsed_response[0]
             return json.dumps(section)
-        return f"No route found for ship '{ship}' at '{position}' with '{moneyToSpend}' aUEC."
+        return f"No route found for ship '{ship}' at '{position}' with '{money_to_spend}' aUEC."
 
     def _get_celestial_object_id(self, name: str) -> Optional[int]:
         """Finds the ID of the celestial object with the specified name."""
diff --git a/skills/vision_ai/main.py b/skills/vision_ai/main.py
@@ -66,70 +66,78 @@ async def execute_tool(
         instant_response = ""
 
         if tool_name == "analyse_what_you_or_user_sees":
-            # Take a screenshot
-            with mss() as sct:
-                main_display = sct.monitors[self.display]
-                screenshot = sct.grab(main_display)
-
-                # Create a PIL image from array
-                image = Image.frombytes(
-                    "RGB", screenshot.size, screenshot.bgra, "raw", "BGRX"
-                )
-
-                desired_width = 1000
-                aspect_ratio = image.height / image.width
-                new_height = int(desired_width * aspect_ratio)
-
-                resized_image = image.resize((desired_width, new_height))
 
-                png_base64 = self.pil_image_to_base64(resized_image)
+            question = parameters.get("question", "What's in this image?")
+            answer = await self.analyse_screen(question)
 
-                if self.show_screenshots:
+            if answer:
+                if self.settings.debug_mode:
                     await self.printr.print_async(
-                        "Analyzing this image",
-                        color=LogType.INFO,
-                        source=LogSource.WINGMAN,
-                        source_name=self.wingman.name,
-                        skill_name=self.name,
-                        additional_data={"image_base64": png_base64},
+                        f"Vision analysis: {answer}.", color=LogType.INFO
                     )
+                function_response = answer
 
-                question = parameters.get("question", "What's in this image?")
+        return function_response, instant_response
 
-                messages = [
-                    {
-                        "role": "system",
-                        "content": """
-                            You are a helpful ai assistant.
-                        """,
-                    },
-                    {
-                        "role": "user",
-                        "content": [
-                            {"type": "text", "text": question},
-                            {
-                                "type": "image_url",
-                                "image_url": {
-                                    "url": f"data:image/jpeg;base64,{png_base64}",
-                                    "detail": "high",
-                                },
-                            },
-                        ],
-                    },
-                ]
-                completion = await self.llm_call(messages)
-                answer = (
-                    completion.choices[0].message.content
-                    if completion and completion.choices
-                    else ""
-                )
+    async def analyse_screen(self, prompt: str, desired_image_width: int = 1000):
+        function_response = ""
 
-                if answer:
-                    if self.settings.debug_mode:
-                        await self.printr.print_async(f"Vision analysis: {answer}.", color=LogType.INFO)
-                    function_response = answer
+        # Take a screenshot
+        with mss() as sct:
+            main_display = sct.monitors[self.display]
+            screenshot = sct.grab(main_display)
 
-        return function_response, instant_response
+            # Create a PIL image from array
+            image = Image.frombytes(
+                "RGB", screenshot.size, screenshot.bgra, "raw", "BGRX"
+            )
+
+            aspect_ratio = image.height / image.width
+            new_height = int(desired_image_width * aspect_ratio)
+
+            resized_image = image.resize((desired_image_width, new_height))
+
+            png_base64 = self.pil_image_to_base64(resized_image)
+
+            if self.show_screenshots:
+                await self.printr.print_async(
+                    "Analyzing this image",
+                    color=LogType.INFO,
+                    source=LogSource.WINGMAN,
+                    source_name=self.wingman.name,
+                    skill_name=self.name,
+                    additional_data={"image_base64": png_base64},
+                )
+
+            messages = [
+                {
+                    "role": "system",
+                    "content": """
+                        You are a helpful ai assistant.
+                    """,
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{png_base64}",
+                                "detail": "high",
+                            },
+                        },
+                    ],
+                },
+            ]
+            completion = await self.llm_call(messages)
+            function_response = (
+                completion.choices[0].message.content
+                if completion and completion.choices
+                else ""
+            )
+
+        return function_response
 
     async def is_summarize_needed(self, tool_name: str) -> bool:
         """Returns whether a tool needs to be summarized."""
diff --git a/templates/skills/star_head/default_config.yaml b/templates/skills/star_head/default_config.yaml
diff --git a/templates/skills/star_head/main.py b/templates/skills/star_head/main.py
