synth-ai 0.2.8.dev2__py3-none-any.whl → 0.4.3__py3-none-any.whl

synth_ai/lm/vendors/core/gemini_api.py

@@ -1,292 +0,0 @@
-import json
-import logging
-import os
-import warnings
-from typing import Any
-
-import google.genai as genai
-from google.api_core.exceptions import ResourceExhausted
-from google.genai import types
-
-from synth_ai.lm.caching.initialize import get_cache_handler
-from synth_ai.lm.constants import (
-    GEMINI_REASONING_MODELS,
-    GEMINI_THINKING_BUDGETS,
-    SPECIAL_BASE_TEMPS,
-)
-from synth_ai.lm.tools.base import BaseTool
-from synth_ai.lm.vendors.base import BaseLMResponse, VendorBase
-from synth_ai.lm.vendors.retries import BACKOFF_TOLERANCE, MAX_BACKOFF, backoff
-
-ALIASES = {
-    "gemini-2.5-flash": "gemini-2.5-flash-preview-04-17",
-}
-
-logger = logging.getLogger(__name__)
-_CLIENT = None  # Initialize lazily when needed
-GEMINI_EXCEPTIONS_TO_RETRY: tuple[type[Exception], ...] = (ResourceExhausted,)
-logging.getLogger("google.genai").setLevel(logging.ERROR)
-os.environ["GRPC_VERBOSITY"] = "ERROR"
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-warnings.filterwarnings("ignore")
-
-SAFETY_SETTINGS = {
-    types.HarmCategory.HARM_CATEGORY_HATE_SPEECH: types.HarmBlockThreshold.BLOCK_NONE,
-    types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: types.HarmBlockThreshold.BLOCK_NONE,
-    types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: types.HarmBlockThreshold.BLOCK_NONE,
-    types.HarmCategory.HARM_CATEGORY_HARASSMENT: types.HarmBlockThreshold.BLOCK_NONE,
-}
-
-
-def _get_client():
-    """Get or create the Gemini client lazily."""
-    global _CLIENT
-    if _CLIENT is None:
-        _CLIENT = genai.Client()
-    return _CLIENT
-
-
-class GeminiAPI(VendorBase):
-    used_for_structured_outputs: bool = True
-    exceptions_to_retry: tuple[type[Exception], ...] = GEMINI_EXCEPTIONS_TO_RETRY
-
-    def __init__(
-        self,
-        exceptions_to_retry: tuple[type[Exception], ...] = GEMINI_EXCEPTIONS_TO_RETRY,
-        used_for_structured_outputs: bool = False,
-    ):
-        self.used_for_structured_outputs = used_for_structured_outputs
-        self.exceptions_to_retry = exceptions_to_retry
-
-    def get_aliased_model_name(self, model_name: str) -> str:
-        if model_name in ALIASES:
-            return ALIASES[model_name]
-        return model_name
-
-    @staticmethod
-    def _msg_to_contents(messages: list[dict[str, Any]]) -> list[types.Content]:
-        # contents, sys_instr = [], None
-        contents = []
-        for m in messages:
-            # if m["role"] == "system":
-            #     sys_instr = f"<instructions>\n{m['content']}\n</instructions>"
-            #     continue
-            # text = (sys_instr + "\n" + m["content"]) if sys_instr else m["content"]
-            if m["role"].lower() not in ["user", "assistant"]:
-                continue
-            role = "user" if m["role"] == "user" else "assistant"
-            contents.append(
-                types.Content(role=role, parts=[types.Part.from_text(text=m["content"])])
-            )
-        return contents
-
-    @staticmethod
-    def _tools_to_genai(tools: list[BaseTool]) -> list[types.Tool]:
-        """Convert internal BaseTool → genai Tool."""
-        out: list[types.Tool] = []
-        for t in tools:
-            # Assume t.to_gemini_tool() now correctly returns a FunctionDeclaration
-            func_decl = t if isinstance(t, dict) else t.to_gemini_tool()
-            if not isinstance(func_decl, types.FunctionDeclaration):
-                # Or fetch schema parts if to_gemini_tool still returns dict
-                # This depends on BaseTool.to_gemini_tool implementation
-                tool_dict = func_decl  # Assuming it's a dict for now
-                func_decl = types.FunctionDeclaration(
-                    name=tool_dict["name"],
-                    description=tool_dict["description"],
-                    parameters=tool_dict["parameters"],  # Expects OpenAPI-style dict
-                )
-            out.append(types.Tool(function_declarations=[func_decl]))
-        return out
-
-    async def _gen_content_async(
-        self,
-        messages: list[dict],
-        temperature: float,
-        model_name: str,
-        reasoning_effort: str,
-        tools: list[BaseTool] | None,
-        lm_config: dict[str, Any] | None,
-    ) -> tuple[str, list[dict] | None]:
-        model_name = self.get_aliased_model_name(model_name)
-        cfg_kwargs: dict[str, Any] = {"temperature": temperature}
-        if model_name in GEMINI_REASONING_MODELS and reasoning_effort in GEMINI_THINKING_BUDGETS:
-            cfg_kwargs["thinking_config"] = types.ThinkingConfig(
-                thinking_budget=GEMINI_THINKING_BUDGETS[reasoning_effort]
-            )
-
-        if any(m["role"] == "system" for m in messages):
-            cfg_kwargs["system_instruction"] = next(
-                m["content"] for m in messages if m["role"] == "system"
-            )
-
-        generation_config = types.GenerateContentConfig(
-            **cfg_kwargs,
-            tool_config=lm_config.get("tool_config") if lm_config else None,
-            tools=self._tools_to_genai(tools) if tools else None,
-        )
-        client = _get_client()
-        resp = await client.aio.models.generate_content(
-            model=model_name,
-            contents=self._msg_to_contents(messages),
-            config=generation_config,
-            # safety_settings=SAFETY_SETTINGS,
-        )
-        return self._extract(resp)
-
-    def _gen_content_sync(
-        self,
-        messages: list[dict],
-        temperature: float,
-        model_name: str,
-        reasoning_effort: str,
-        tools: list[BaseTool] | None,
-        lm_config: dict[str, Any] | None,
-    ) -> tuple[str, list[dict] | None]:
-        model_name = self.get_aliased_model_name(model_name)
-        cfg_kwargs: dict[str, Any] = {"temperature": temperature}
-        if model_name in GEMINI_REASONING_MODELS and reasoning_effort in GEMINI_THINKING_BUDGETS:
-            cfg_kwargs["thinking_config"] = types.ThinkingConfig(
-                thinking_budget=GEMINI_THINKING_BUDGETS[reasoning_effort]
-            )
-        if any(m["role"] == "system" for m in messages):
-            cfg_kwargs["system_instruction"] = next(
-                m["content"] for m in messages if m["role"] == "system"
-            )
-        generation_config = types.GenerateContentConfig(
-            **cfg_kwargs,
-            tool_config=lm_config.get("tool_config") if lm_config else None,
-            tools=self._tools_to_genai(tools) if tools else None,
-        )
-
-        client = _get_client()
-        resp = client.models.generate_content(
-            model=model_name,
-            contents=self._msg_to_contents(messages),
-            safety_settings=SAFETY_SETTINGS,
-            config=generation_config,
-        )
-        return self._extract(resp)
-
-    @staticmethod
-    def _extract(response) -> tuple[str, list[dict] | None]:
-        # Extract text, handling cases where it might be missing
-        try:
-            text = response.text
-        except ValueError:  # Handle cases where only non-text parts exist
-            text = ""
-
-        calls = []
-        # Access parts through candidates[0].content
-        if response.candidates and response.candidates[0].content:
-            for part in response.candidates[0].content.parts:
-                if part.function_call:
-                    calls.append(
-                        {
-                            "id": f"call_{len(calls) + 1}",
-                            "type": "function",
-                            "function": {
-                                "name": part.function_call.name,
-                                "arguments": json.dumps(dict(part.function_call.args)),
-                            },
-                        }
-                    )
-        return text, calls or None
-
-    @backoff.on_exception(
-        backoff.expo,
-        exceptions_to_retry,
-        max_tries=BACKOFF_TOLERANCE,
-        max_time=MAX_BACKOFF,
-        on_giveup=lambda e: print(e),
-    )
-    async def _hit_api_async(
-        self,
-        model: str,
-        messages: list[dict[str, Any]],
-        lm_config: dict[str, Any],
-        use_ephemeral_cache_only: bool = False,
-        reasoning_effort: str = "high",
-        tools: list[BaseTool] | None = None,
-    ) -> BaseLMResponse:
-        assert lm_config.get("response_model") is None, (
-            "response_model is not supported for standard calls"
-        )
-        used_cache_handler = get_cache_handler(use_ephemeral_cache_only)
-        lm_config["reasoning_effort"] = reasoning_effort
-        cache_result = used_cache_handler.hit_managed_cache(
-            model, messages, lm_config=lm_config, tools=tools
-        )
-        if cache_result:
-            return cache_result
-
-        raw_response, tool_calls = await self._gen_content_async(
-            messages,
-            temperature=lm_config.get("temperature", SPECIAL_BASE_TEMPS.get(model, 0)),
-            reasoning_effort=reasoning_effort,
-            tools=tools,
-            lm_config=lm_config,
-            model_name=model,
-        )
-        if not raw_response:
-            raw_response = ""
-        lm_response = BaseLMResponse(
-            raw_response=raw_response,
-            structured_output=None,
-            tool_calls=tool_calls,
-        )
-
-        lm_config["reasoning_effort"] = reasoning_effort
-        used_cache_handler.add_to_managed_cache(
-            model, messages, lm_config=lm_config, output=lm_response, tools=tools
-        )
-        return lm_response
-
-    @backoff.on_exception(
-        backoff.expo,
-        exceptions_to_retry,
-        max_tries=BACKOFF_TOLERANCE,
-        max_time=MAX_BACKOFF,
-        on_giveup=lambda e: print(e),
-    )
-    def _hit_api_sync(
-        self,
-        model: str,
-        messages: list[dict[str, Any]],
-        lm_config: dict[str, Any],
-        use_ephemeral_cache_only: bool = False,
-        reasoning_effort: str = "high",
-        tools: list[BaseTool] | None = None,
-    ) -> BaseLMResponse:
-        assert lm_config.get("response_model") is None, (
-            "response_model is not supported for standard calls"
-        )
-        used_cache_handler = get_cache_handler(use_ephemeral_cache_only=use_ephemeral_cache_only)
-        lm_config["reasoning_effort"] = reasoning_effort
-        cache_result = used_cache_handler.hit_managed_cache(
-            model, messages, lm_config=lm_config, tools=tools
-        )
-        if cache_result:
-            return cache_result
-
-        raw_response, tool_calls = self._gen_content_sync(
-            messages,
-            temperature=lm_config.get("temperature", SPECIAL_BASE_TEMPS.get(model, 0)),
-            reasoning_effort=reasoning_effort,
-            tools=tools,
-            lm_config=lm_config,
-            model_name=model,
-        )
-        if not raw_response:
-            raw_response = ""
-        lm_response = BaseLMResponse(
-            raw_response=raw_response,
-            structured_output=None,
-            tool_calls=tool_calls,
-        )
-
-        lm_config["reasoning_effort"] = reasoning_effort
-        used_cache_handler.add_to_managed_cache(
-            model, messages, lm_config=lm_config, output=lm_response, tools=tools
-        )
-        return lm_response

synth_ai/lm/vendors/core/mistral_api.py

@@ -1,322 +0,0 @@
-import json
-import os
-from typing import Any
-
-import pydantic
-from mistralai import Mistral  # use Mistral as both sync and async client
-from pydantic import BaseModel
-
-from synth_ai.lm.caching.initialize import get_cache_handler
-from synth_ai.lm.constants import SPECIAL_BASE_TEMPS
-from synth_ai.lm.tools.base import BaseTool
-from synth_ai.lm.vendors.base import BaseLMResponse, VendorBase
-from synth_ai.lm.vendors.core.openai_api import OpenAIStructuredOutputClient
-
-# Since the mistralai package doesn't expose an exceptions module,
-# we fallback to catching all Exceptions for retry.
-MISTRAL_EXCEPTIONS_TO_RETRY: tuple[type[Exception], ...] = (Exception,)
-
-
-class MistralAPI(VendorBase):
-    used_for_structured_outputs: bool = True
-    exceptions_to_retry: tuple = MISTRAL_EXCEPTIONS_TO_RETRY
-    _openai_fallback: Any
-
-    def __init__(
-        self,
-        exceptions_to_retry: tuple[type[Exception], ...] = MISTRAL_EXCEPTIONS_TO_RETRY,
-        used_for_structured_outputs: bool = False,
-    ):
-        self.used_for_structured_outputs = used_for_structured_outputs
-        self.exceptions_to_retry = exceptions_to_retry
-        self._openai_fallback = None
-
-    # @backoff.on_exception(
-    #     backoff.expo,
-    #     MISTRAL_EXCEPTIONS_TO_RETRY,
-    #     max_tries=BACKOFF_TOLERANCE,
-    #     on_giveup=lambda e: print(e),
-    # )
-    async def _hit_api_async(
-        self,
-        model: str,
-        messages: list[dict[str, Any]],
-        lm_config: dict[str, Any],
-        response_model: BaseModel | None = None,
-        use_ephemeral_cache_only: bool = False,
-        reasoning_effort: str = "high",
-        tools: list[BaseTool] | None = None,
-    ) -> BaseLMResponse:
-        assert lm_config.get("response_model") is None, (
-            "response_model is not supported for standard calls"
-        )
-        assert not (response_model and tools), "Cannot provide both response_model and tools"
-        used_cache_handler = get_cache_handler(use_ephemeral_cache_only)
-        lm_config["reasoning_effort"] = reasoning_effort
-        cache_result = used_cache_handler.hit_managed_cache(
-            model, messages, lm_config=lm_config, tools=tools
-        )
-        if cache_result:
-            assert type(cache_result) in [
-                BaseLMResponse,
-                str,
-            ], f"Expected BaseLMResponse or str, got {type(cache_result)}"
-            return (
-                cache_result
-                if isinstance(cache_result, BaseLMResponse)
-                else BaseLMResponse(
-                    raw_response=cache_result, structured_output=None, tool_calls=None
-                )
-            )
-
-        mistral_messages = [{"role": msg["role"], "content": msg["content"]} for msg in messages]
-        functions = [tool.to_mistral_tool() for tool in tools] if tools else None
-        params = {
-            "model": model,
-            "messages": mistral_messages,
-            "max_tokens": lm_config.get("max_tokens", 4096),
-            "temperature": lm_config.get("temperature", SPECIAL_BASE_TEMPS.get(model, 0)),
-            "stream": False,
-            "tool_choice": "auto" if functions else None,
-        }
-        if response_model:
-            params["response_format"] = response_model
-        elif tools:
-            params["tools"] = functions
-
-        async with Mistral(api_key=os.getenv("MISTRAL_API_KEY", "")) as client:
-            response = await client.chat.complete_async(**params)
-
-        message = response.choices[0].message
-        try:
-            raw_response = message.content
-        except AttributeError:
-            raw_response = ""
-
-        tool_calls = []
-        try:
-            if message.tool_calls:
-                tool_calls = [
-                    {
-                        "id": call.id,
-                        "type": "function",
-                        "function": {
-                            "name": call.function.name,
-                            "arguments": call.function.arguments,
-                        },
-                    }
-                    for call in message.tool_calls
-                ]
-        except AttributeError:
-            pass
-
-        lm_response = BaseLMResponse(
-            raw_response=raw_response,
-            structured_output=None,
-            tool_calls=tool_calls if tool_calls else None,
-        )
-        lm_config["reasoning_effort"] = reasoning_effort
-        used_cache_handler.add_to_managed_cache(
-            model, messages, lm_config=lm_config, output=lm_response, tools=tools
-        )
-        return lm_response
-
-    # @backoff.on_exception(
-    #     backoff.expo,
-    #     MISTRAL_EXCEPTIONS_TO_RETRY,
-    #     max_tries=BACKOFF_TOLERANCE,
-    #     on_giveup=lambda e: print(e),
-    # )
-    def _hit_api_sync(
-        self,
-        model: str,
-        messages: list[dict[str, Any]],
-        lm_config: dict[str, Any],
-        response_model: BaseModel | None = None,
-        use_ephemeral_cache_only: bool = False,
-        reasoning_effort: str = "high",
-        tools: list[BaseTool] | None = None,
-    ) -> BaseLMResponse:
-        assert lm_config.get("response_model") is None, (
-            "response_model is not supported for standard calls"
-        )
-        assert not (response_model and tools), "Cannot provide both response_model and tools"
-
-        used_cache_handler = get_cache_handler(use_ephemeral_cache_only)
-        lm_config["reasoning_effort"] = reasoning_effort
-        cache_result = used_cache_handler.hit_managed_cache(
-            model, messages, lm_config=lm_config, tools=tools
-        )
-        if cache_result:
-            assert type(cache_result) in [
-                BaseLMResponse,
-                str,
-            ], f"Expected BaseLMResponse or str, got {type(cache_result)}"
-            return (
-                cache_result
-                if isinstance(cache_result, BaseLMResponse)
-                else BaseLMResponse(
-                    raw_response=cache_result, structured_output=None, tool_calls=None
-                )
-            )
-
-        mistral_messages = [{"role": msg["role"], "content": msg["content"]} for msg in messages]
-        functions = [tool.to_mistral_tool() for tool in tools] if tools else None
-
-        params = {
-            "model": model,
-            "messages": mistral_messages,
-            "max_tokens": lm_config.get("max_tokens", 4096),
-            "temperature": lm_config.get("temperature", SPECIAL_BASE_TEMPS.get(model, 0)),
-            "stream": False,
-            "tool_choice": "auto" if functions else None,
-            # "tools": functions,
-        }
-        if response_model:
-            params["response_format"] = response_model
-        elif tools:
-            params["tools"] = functions
-
-        with Mistral(api_key=os.getenv("MISTRAL_API_KEY", "")) as client:
-            response = client.chat.complete(**params)
-
-        message = response.choices[0].message
-        try:
-            raw_response = message.content
-        except AttributeError:
-            raw_response = ""
-
-        tool_calls = []
-        try:
-            if message.tool_calls:
-                tool_calls = [
-                    {
-                        "id": call.id,
-                        "type": "function",
-                        "function": {
-                            "name": call.function.name,
-                            "arguments": call.function.arguments,
-                        },
-                    }
-                    for call in message.tool_calls
-                ]
-        except AttributeError:
-            pass
-
-        lm_response = BaseLMResponse(
-            raw_response=raw_response,
-            structured_output=None,
-            tool_calls=tool_calls if tool_calls else None,
-        )
-        lm_config["reasoning_effort"] = reasoning_effort
-        used_cache_handler.add_to_managed_cache(
-            model, messages, lm_config=lm_config, output=lm_response, tools=tools
-        )
-        return lm_response
-
-    async def _hit_api_async_structured_output(
-        self,
-        model: str,
-        messages: list[dict[str, Any]],
-        response_model: BaseModel,
-        temperature: float,
-        use_ephemeral_cache_only: bool = False,
-    ) -> BaseLMResponse:
-        try:
-            mistral_messages = [
-                {"role": msg["role"], "content": msg["content"]} for msg in messages
-            ]
-            async with Mistral(api_key=os.getenv("MISTRAL_API_KEY", "")) as client:
-                response = await client.chat.complete_async(
-                    model=model,
-                    messages=mistral_messages,
-                    max_tokens=4096,
-                    temperature=temperature,
-                    stream=False,
-                )
-            result = response.choices[0].message.content
-            parsed = json.loads(result)
-            lm_response = BaseLMResponse(
-                raw_response="",
-                structured_output=response_model(**parsed),
-                tool_calls=None,
-            )
-            return lm_response
-        except (json.JSONDecodeError, pydantic.ValidationError):
-            if self._openai_fallback is None:
-                self._openai_fallback = OpenAIStructuredOutputClient()
-            return await self._openai_fallback._hit_api_async_structured_output(
-                model="gpt-4o",
-                messages=messages,
-                response_model=response_model,
-                temperature=temperature,
-                use_ephemeral_cache_only=use_ephemeral_cache_only,
-            )
-
-    def _hit_api_sync_structured_output(
-        self,
-        model: str,
-        messages: list[dict[str, Any]],
-        response_model: BaseModel,
-        temperature: float,
-        use_ephemeral_cache_only: bool = False,
-    ) -> BaseLMResponse:
-        try:
-            mistral_messages = [
-                {"role": msg["role"], "content": msg["content"]} for msg in messages
-            ]
-            with Mistral(api_key=os.getenv("MISTRAL_API_KEY", "")) as client:
-                response = client.chat.complete(
-                    model=model,
-                    messages=mistral_messages,
-                    max_tokens=4096,
-                    temperature=temperature,
-                    stream=False,
-                )
-            result = response.choices[0].message.content
-            parsed = json.loads(result)
-            lm_response = BaseLMResponse(
-                raw_response="",
-                structured_output=response_model(**parsed),
-                tool_calls=None,
-            )
-            return lm_response
-        except (json.JSONDecodeError, pydantic.ValidationError):
-            print("WARNING - Falling back to OpenAI - THIS IS SLOW")
-            if self._openai_fallback is None:
-                self._openai_fallback = OpenAIStructuredOutputClient()
-            return self._openai_fallback._hit_api_sync_structured_output(
-                model="gpt-4o",
-                messages=messages,
-                response_model=response_model,
-                temperature=temperature,
-                use_ephemeral_cache_only=use_ephemeral_cache_only,
-            )
-
-
-if __name__ == "__main__":
-    import asyncio
-
-    from pydantic import BaseModel
-
-    class TestModel(BaseModel):
-        name: str
-
-    client = MistralAPI(used_for_structured_outputs=True, exceptions_to_retry=[])
-    import time
-
-    t = time.time()
-
-    async def run_async():
-        response = await client._hit_api_async_structured_output(
-            model="mistral-large-latest",
-            messages=[{"role": "user", "content": "What is the capital of the moon?"}],
-            response_model=TestModel,
-            temperature=0.0,
-        )
-        print(response)
-        return response
-
-    response = asyncio.run(run_async())
-    t2 = time.time()
-    print(f"Got {len(response.name)} chars in {t2 - t} seconds")