synth-ai 0.1.0.dev28__py3-none-any.whl → 0.1.0.dev30__py3-none-any.whl

synth_ai/zyk/lms/core/main.py

@@ -9,7 +9,8 @@ from synth_ai.zyk.lms.core.vendor_clients import (
     openai_naming_regexes,
 )
 from synth_ai.zyk.lms.structured_outputs.handler import StructuredOutputHandler
-from synth_ai.zyk.lms.vendors.base import VendorBase
+from synth_ai.zyk.lms.vendors.base import VendorBase, BaseLMResponse
+from synth_ai.zyk.lms.tools.base import BaseTool
 
 REASONING_MODELS = ["deepseek-reasoner", "o1-mini", "o1-preview", "o1", "o3"]
 
@@ -120,22 +121,23 @@ class LM:
         images_as_bytes: List[Any] = [],
         response_model: Optional[BaseModel] = None,
         use_ephemeral_cache_only: bool = False,
-    ):
+        tools: Optional[List[BaseTool]] = None,
+    ) -> BaseLMResponse:
         assert (system_message is None) == (
             user_message is None
         ), "Must provide both system_message and user_message or neither"
         assert (
             (messages is None) != (system_message is None)
         ), "Must provide either messages or system_message/user_message pair, but not both"
-
+        assert not (response_model and tools), "Cannot provide both response_model and tools"
         if messages is None:
             messages = build_messages(
                 system_message, user_message, images_as_bytes, self.model_name
             )
-
+        result = None
         if response_model:
             try:
-                return self.structured_output_handler.call_sync(
+                result = self.structured_output_handler.call_sync(
                     messages,
                     model=self.model_name,
                     lm_config=self.lm_config,
@@ -144,7 +146,7 @@ class LM:
                 )
             except StructuredOutputCoercionFailureException:
                 # print("Falling back to backup handler")
-                return self.backup_structured_output_handler.call_sync(
+                result = self.backup_structured_output_handler.call_sync(
                     messages,
                     model=self.model_name,
                     lm_config=self.lm_config,
@@ -152,12 +154,17 @@ class LM:
                     use_ephemeral_cache_only=use_ephemeral_cache_only,
                 )
         else:
-            return self.client._hit_api_sync(
+            result = self.client._hit_api_sync(
                 messages=messages,
                 model=self.model_name,
                 lm_config=self.lm_config,
                 use_ephemeral_cache_only=use_ephemeral_cache_only,
+                tools=tools,
             )
+        assert isinstance(result.raw_response, str), "Raw response must be a string"
+        assert (isinstance(result.structured_output, BaseModel) or result.structured_output is None), "Structured output must be a Pydantic model or None"
+        assert (isinstance(result.tool_calls, list) or result.tool_calls is None), "Tool calls must be a list or None"
+        return result
 
     async def respond_async(
         self,
@@ -167,7 +174,8 @@ class LM:
         images_as_bytes: List[Any] = [],
         response_model: Optional[BaseModel] = None,
         use_ephemeral_cache_only: bool = False,
-    ):
+        tools: Optional[List[BaseTool]] = None,
+    ) -> BaseLMResponse:
         # "In respond_async")
         assert (system_message is None) == (
             user_message is None
@@ -176,15 +184,16 @@ class LM:
             (messages is None) != (system_message is None)
         ), "Must provide either messages or system_message/user_message pair, but not both"
 
+        assert not (response_model and tools), "Cannot provide both response_model and tools"
         if messages is None:
             messages = build_messages(
                 system_message, user_message, images_as_bytes, self.model_name
             )
-
+        result = None
         if response_model:
             try:
-                # "Trying structured output handler")
-                return await self.structured_output_handler.call_async(
+                #print("Trying structured output handler")
+                result = await self.structured_output_handler.call_async(
                     messages,
                     model=self.model_name,
                     lm_config=self.lm_config,
@@ -192,8 +201,8 @@ class LM:
                     use_ephemeral_cache_only=use_ephemeral_cache_only,
                 )
             except StructuredOutputCoercionFailureException:
-                # print("Falling back to backup handler")
-                return await self.backup_structured_output_handler.call_async(
+                #print("Falling back to backup handler")
+                result = await self.backup_structured_output_handler.call_async(
                     messages,
                     model=self.model_name,
                     lm_config=self.lm_config,
@@ -201,14 +210,18 @@ class LM:
                     use_ephemeral_cache_only=use_ephemeral_cache_only,
                 )
         else:
-            # print("Calling API no response model")
-            return await self.client._hit_api_async(
+            #print("Calling API no response model")
+            result = await self.client._hit_api_async(
                 messages=messages,
                 model=self.model_name,
                 lm_config=self.lm_config,
                 use_ephemeral_cache_only=use_ephemeral_cache_only,
+                tools=tools,
             )
-
+        assert isinstance(result.raw_response, str), "Raw response must be a string"
+        assert (isinstance(result.structured_output, BaseModel) or result.structured_output is None), "Structured output must be a Pydantic model or None"
+        assert (isinstance(result.tool_calls, list) or result.tool_calls is None), "Tool calls must be a list or None"
+        return result
 
 if __name__ == "__main__":
     import asyncio

synth_ai/zyk/lms/core/vendor_clients.py

@@ -6,8 +6,8 @@ from synth_ai.zyk.lms.core.all import (
     DeepSeekClient,
     GeminiClient,
     GroqAPI,
-    # OpenAIClient,
     MistralAPI,
+    # OpenAIClient,
     OpenAIStructuredOutputClient,
     TogetherClient,
 )

synth_ai/zyk/lms/structured_outputs/handler.py

@@ -1,3 +1,4 @@
+import logging
 import time
 from abc import ABC, abstractmethod
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
@@ -11,13 +12,13 @@ from synth_ai.zyk.lms.structured_outputs.inject import (
 from synth_ai.zyk.lms.structured_outputs.rehabilitate import (
     fix_errant_forced_async,
     fix_errant_forced_sync,
-    fix_errant_stringified_json_async,
-    fix_errant_stringified_json_sync,
     pull_out_structured_output,
 )
-from synth_ai.zyk.lms.vendors.base import VendorBase
+from synth_ai.zyk.lms.vendors.base import BaseLMResponse, VendorBase
 from synth_ai.zyk.lms.vendors.constants import SPECIAL_BASE_TEMPS
 
+logger = logging.getLogger(__name__)
+
 
 class StructuredHandlerBase(ABC):
     core_client: VendorBase
@@ -49,7 +50,7 @@ class StructuredHandlerBase(ABC):
         temperature: float = 0.0,
         use_ephemeral_cache_only: bool = False,
         reasoning_effort: str = "high",
-    ) -> BaseModel:
+    ) -> BaseLMResponse:
         if temperature == 0.0:
             temperature = SPECIAL_BASE_TEMPS.get(model, 0.0)
         # print("Calling from base")
@@ -73,7 +74,7 @@ class StructuredHandlerBase(ABC):
         temperature: float = 0.0,
         use_ephemeral_cache_only: bool = False,
         reasoning_effort: str = "high",
-    ) -> BaseModel:
+    ) -> BaseLMResponse:
         if temperature == 0.0:
             temperature = SPECIAL_BASE_TEMPS.get(model, 0.0)
         return self._process_call_sync(
@@ -97,7 +98,7 @@ class StructuredHandlerBase(ABC):
         api_call_method,
         use_ephemeral_cache_only: bool = False,
         reasoning_effort: str = "high",
-    ) -> BaseModel:
+    ) -> BaseLMResponse:
         pass
 
     @abstractmethod
@@ -109,7 +110,7 @@ class StructuredHandlerBase(ABC):
         api_call_method,
         use_ephemeral_cache_only: bool = False,
         reasoning_effort: str = "high",
-    ) -> BaseModel:
+    ) -> BaseLMResponse:
         pass
 
 
@@ -140,11 +141,9 @@ class StringifiedJSONHandler(StructuredHandlerBase):
         api_call_method: Callable,
         use_ephemeral_cache_only: bool = False,
         reasoning_effort: str = "high",
-    ) -> BaseModel:
-        # print("In _process_call_async")
-        assert isinstance(
-            api_call_method, Callable
-        ), "api_call_method must be a callable"
+    ) -> BaseLMResponse:
+        logger.info(f"Processing structured output call for model: {model}")
+        assert callable(api_call_method), "api_call_method must be a callable"
         assert (
             response_model is not None
         ), "Don't use this handler for unstructured outputs"
@@ -168,43 +167,63 @@ class StringifiedJSONHandler(StructuredHandlerBase):
                 use_ephemeral_cache_only=use_ephemeral_cache_only,
                 reasoning_effort=reasoning_effort,
             )
-            # print(f"Time to get response: {time.time() - t0}")
-            if not isinstance(raw_text_response_or_cached_hit, str):
-                return raw_text_response_or_cached_hit
+            logger.debug(f"Time to get response: {time.time() - t0:.2f}s")
+
+            # Check if we got a cached BaseLMResponse
+            assert (
+                type(raw_text_response_or_cached_hit) in [str, BaseLMResponse]
+            ), f"Expected str or BaseLMResponse, got {type(raw_text_response_or_cached_hit)}"
+            if type(raw_text_response_or_cached_hit) == BaseLMResponse:
+                print("Got cached hit, returning directly")
+                raw_text_response = raw_text_response_or_cached_hit.raw_response
             else:
                 raw_text_response = raw_text_response_or_cached_hit
+            logger.debug(f"Raw response from model:\n{raw_text_response}")
+
+            print("Trying to parse structured output")
             try:
                 structured_output = pull_out_structured_output(
                     raw_text_response, response_model
                 )
+
+                print("Successfully parsed structured output on first attempt")
                 break
-            # except Exception as e:
-            #     try:
-            #         structured_output = await fix_errant_stringified_json_async(raw_text_response, response_model)
-            #         break
             except Exception as e:
+                logger.warning(f"Failed to parse structured output: {str(e)}")
                 try:
-                    # t0 = time.time()
-                    # print(f"Got error {e}, attempting to fix")
+                    print("Attempting to fix with forced JSON parser")
                     structured_output = await fix_errant_forced_async(
                         messages_with_json_formatting_instructions,
                         raw_text_response,
                         response_model,
                         "gpt-4o-mini",
                     )
-
-                    # print(f"Time to fix: {time.time() - t0}")
+                    assert isinstance(structured_output, BaseModel), "Structured output must be a Pydantic model"
+                    assert not isinstance(structured_output, BaseLMResponse), "Got BaseLMResponse instead of Pydantic model"
+                    print("Successfully fixed and parsed structured output")
                     break
                 except Exception as e:
+                    logger.error(f"Failed to fix structured output: {str(e)}")
                     previously_failed_error_messages.append(
                         f"Generated attempt and got error. Attempt:\n\n{raw_text_response}\n\nError:\n\n{e}"
                     )
                     remaining_retries -= 1
+                    logger.warning(f"Retries remaining: {remaining_retries}")
+
         if structured_output is None:
+            logger.error("Failed to get structured output after all retries")
             raise StructuredOutputCoercionFailureException(
                 "Failed to get structured output"
             )
-        return structured_output
+        print("Successfully parsed structured output")
+        print(structured_output)
+        assert isinstance(structured_output, BaseModel), "Structured output must be a Pydantic model"
+        assert not isinstance(structured_output, BaseLMResponse),"Got BaseLMResponse instead of Pydantic model" 
+        return BaseLMResponse(
+            raw_response=raw_text_response,
+            structured_output=structured_output,
+            tool_calls=None,
+        )
 
     def _process_call_sync(
         self,
@@ -215,10 +234,9 @@ class StringifiedJSONHandler(StructuredHandlerBase):
         api_call_method: Callable,
         use_ephemeral_cache_only: bool = False,
         reasoning_effort: str = "high",
-    ) -> BaseModel:
-        assert isinstance(
-            api_call_method, Callable
-        ), "api_call_method must be a callable"
+    ) -> BaseLMResponse:
+        logger.info(f"Processing structured output call for model: {model}")
+        assert callable(api_call_method), "api_call_method must be a callable"
         assert (
             response_model is not None
         ), "Don't use this handler for unstructured outputs"
@@ -234,7 +252,7 @@ class StringifiedJSONHandler(StructuredHandlerBase):
                     previously_failed_error_messages=previously_failed_error_messages,
                 )
             )
-            # t0 = time.time()
+            t0 = time.time()
             raw_text_response_or_cached_hit = api_call_method(
                 messages=messages_with_json_formatting_instructions,
                 model=model,
@@ -242,39 +260,54 @@ class StringifiedJSONHandler(StructuredHandlerBase):
                 use_ephemeral_cache_only=use_ephemeral_cache_only,
                 reasoning_effort=reasoning_effort,
             )
-            # print(f"Time to get response: {time.time() - t0}")
-            if not isinstance(raw_text_response_or_cached_hit, str):
-                return raw_text_response_or_cached_hit
-            else:
+            logger.debug(f"Time to get response: {time.time() - t0:.2f}s")
+
+            # Check if we got a cached BaseLMResponse
+            assert (
+                type(raw_text_response_or_cached_hit) in [str, BaseLMResponse]
+            ), f"Expected str or BaseLMResponse, got {type(raw_text_response_or_cached_hit)}"
+            if type(raw_text_response_or_cached_hit) == BaseLMResponse:
+                logger.info("Got cached hit, returning directly")
+                raw_text_response = raw_text_response_or_cached_hit.raw_response
+            else:   
                 raw_text_response = raw_text_response_or_cached_hit
+            logger.debug(f"Raw response from model:\n{raw_text_response}")
+
             try:
                 structured_output = pull_out_structured_output(
                     raw_text_response, response_model
                 )
+                print("Successfully parsed structured output on first attempt")
                 break
-            # except Exception:
-            #     try:
-            #         structured_output = fix_errant_stringified_json_sync(raw_text_response, response_model)
-            #         break
             except Exception as e:
+                logger.warning(f"Failed to parse structured output: {str(e)}")
                 try:
-                    # t0 = time.time()
-                    # print(f"Got error {e}, attempting to fix")
+                    print("Attempting to fix with forced JSON parser")
                     structured_output = fix_errant_forced_sync(
                         raw_text_response, response_model, "gpt-4o-mini"
                     )
+                    print("Successfully fixed and parsed structured output")
                     break
-                    # print(f"Time to fix: {time.time() - t0}")
                 except Exception as e:
+                    logger.error(f"Failed to fix structured output: {str(e)}")
                     previously_failed_error_messages.append(
                         f"Generated attempt and got error. Attempt:\n\n{raw_text_response}\n\nError:\n\n{e}"
                     )
                     remaining_retries -= 1
+                    logger.warning(f"Retries remaining: {remaining_retries}")
+
+        print("Successfully parsed structured output")
+        print(structured_output)
         if structured_output is None:
+            logger.error("Failed to get structured output after all retries")
             raise StructuredOutputCoercionFailureException(
                 "Failed to get structured output"
             )
-        return structured_output
+        return BaseLMResponse(
+            raw_response=raw_text_response,
+            structured_output=structured_output,
+            tool_calls=None,
+        )
 
 
 class ForcedJSONHandler(StructuredHandlerBase):
@@ -296,6 +329,7 @@ class ForcedJSONHandler(StructuredHandlerBase):
             structured_output_mode="forced_json",
         )
         self.reasoning_effort = reasoning_effort
+
     async def _process_call_async(
         self,
         messages: List[Dict[str, Any]],
@@ -305,7 +339,7 @@ class ForcedJSONHandler(StructuredHandlerBase):
         temperature: float = 0.0,
         use_ephemeral_cache_only: bool = False,
         reasoning_effort: str = "high",
-    ) -> BaseModel:
+    ) -> BaseLMResponse:
         # print("Forced JSON")
         assert (
             response_model is not None
@@ -328,7 +362,7 @@ class ForcedJSONHandler(StructuredHandlerBase):
         temperature: float = 0.0,
         use_ephemeral_cache_only: bool = False,
         reasoning_effort: str = "high",
-        ) -> BaseModel:
+    ) -> BaseLMResponse:
         assert (
             response_model is not None
         ), "Don't use this handler for unstructured outputs"
@@ -373,7 +407,7 @@ class StructuredOutputHandler:
         use_ephemeral_cache_only: bool = False,
         lm_config: Dict[str, Any] = {},
         reasoning_effort: str = "high",
-    ) -> BaseModel:
+    ) -> BaseLMResponse:
         # print("Output handler call async")
         return await self.handler.call_async(
             messages=messages,
@@ -394,7 +428,7 @@ class StructuredOutputHandler:
         use_ephemeral_cache_only: bool = False,
         lm_config: Dict[str, Any] = {},
         reasoning_effort: str = "high",
-    ) -> BaseModel:
+    ) -> BaseLMResponse:
         return self.handler.call_sync(
             messages=messages,
             model=model,

synth_ai/zyk/lms/structured_outputs/rehabilitate.py

@@ -55,6 +55,7 @@ def pull_out_structured_output(
         raise ValueError(
             f"Failed to parse response as {response_model}: {e} - {response_prepared}"
         )
+    assert isinstance(final, BaseModel), "Structured output must be a Pydantic model"
     return final
 
 
@@ -157,7 +158,7 @@ async def fix_errant_forced_async(
         )
         # print("Fixed response:")
         # print(fixed_response)
-        return fixed_response
+        return fixed_response.structured_output
 
 
 def fix_errant_forced_sync(
@@ -183,4 +184,4 @@ def fix_errant_forced_sync(
     )
     # print("Fixed response:")
     # print(fixed_response)
-    return fixed_response
+    return fixed_response.structured_output

synth_ai/zyk/lms/tools/base.py

@@ -0,0 +1,104 @@
+from typing import Type
+
+from pydantic import BaseModel
+
+
+class BaseTool(BaseModel):
+    name: str
+    arguments: Type[BaseModel]
+    description: str = ""
+    strict: bool = True
+
+    def to_openai_tool(self):
+        schema = self.arguments.model_json_schema()
+        schema["additionalProperties"] = False
+
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": schema,
+                "strict": self.strict,
+            },
+        }
+
+    def to_anthropic_tool(self):
+        schema = self.arguments.model_json_schema()
+        schema["additionalProperties"] = False
+
+        return {
+            "name": self.name,
+            "description": self.description,
+            "input_schema": {
+                "type": "object",
+                "properties": schema["properties"],
+                "required": schema.get("required", []),
+            },
+        }
+
+    def to_mistral_tool(self):
+        schema = self.arguments.model_json_schema()
+        properties = {}
+        for prop_name, prop in schema.get("properties", {}).items():
+            prop_type = prop["type"]
+            if prop_type == "array" and "items" in prop:
+                properties[prop_name] = {
+                    "type": "array",
+                    "items": prop["items"],
+                    "description": prop.get("description", ""),
+                }
+                continue
+
+            properties[prop_name] = {
+                "type": prop_type,
+                "description": prop.get("description", ""),
+            }
+            if "enum" in prop:
+                properties[prop_name]["enum"] = prop["enum"]
+
+        parameters = {
+            "type": "object",
+            "properties": properties,
+            "required": schema.get("required", []),
+            "additionalProperties": False,
+        }
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": parameters,
+            },
+        }
+
+    def to_gemini_tool(self):
+        schema = self.arguments.model_json_schema()
+        # Convert Pydantic schema types to Gemini schema types
+        properties = {}
+        for name, prop in schema["properties"].items():
+            prop_type = prop.get("type", "string")
+            if prop_type == "array" and "items" in prop:
+                properties[name] = {
+                    "type": "array",
+                    "items": prop["items"],
+                    "description": prop.get("description", ""),
+                }
+                continue
+
+            properties[name] = {
+                "type": prop_type,
+                "description": prop.get("description", ""),
+            }
+            if "enum" in prop:
+                properties[name]["enum"] = prop["enum"]
+
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": {
+                "type": "object",
+                "properties": properties,
+                "required": schema.get("required", []),
+            },
+        }

synth_ai/zyk/lms/vendors/base.py

@@ -1,15 +1,31 @@
 from abc import ABC, abstractmethod
-from typing import List, Dict, Any, Optional
+from typing import Any, Dict, List, Optional
+
 from pydantic import BaseModel
 
+
+class BaseLMResponse(BaseModel):
+    raw_response: str
+    structured_output: Optional[BaseModel] = None
+    tool_calls: Optional[List[Dict]] = None
+
+
 class VendorBase(ABC):
     used_for_structured_outputs: bool = False
     exceptions_to_retry: List[Exception] = []
-    
+
     @abstractmethod
-    async def _hit_api_async(self, messages: List[Dict[str, Any]], response_model_override: Optional[BaseModel] = None) -> str:
+    async def _hit_api_async(
+        self,
+        messages: List[Dict[str, Any]],
+        response_model_override: Optional[BaseModel] = None,
+    ) -> str:
         pass
-    
+
     @abstractmethod
-    def _hit_api_sync(self, messages: List[Dict[str, Any]], response_model_override: Optional[BaseModel] = None) -> str:
-        pass
+    def _hit_api_sync(
+        self,
+        messages: List[Dict[str, Any]],
+        response_model_override: Optional[BaseModel] = None,
+    ) -> str:
+        pass