model-library 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

model_library/providers/xai.py

@@ -1,32 +1,30 @@
-import asyncio
 import io
 import logging
-from typing import Any, Literal, Sequence, cast
+from typing import Any, Literal, Sequence
 
 import grpc
 from typing_extensions import override
-from xai_sdk import AsyncClient, Client
-from xai_sdk.aio.chat import Chat as AsyncChat
+from xai_sdk import AsyncClient
+from xai_sdk.aio.chat import Chat
 from xai_sdk.chat import Content, Response, system, tool_result, user
 from xai_sdk.chat import image as xai_image
 from xai_sdk.chat import tool as xai_tool
 from xai_sdk.proto.v6.chat_pb2 import Message, Tool
-from xai_sdk.sync.chat import Chat as SyncChat
 
 from model_library import model_library_settings
 from model_library.base import (
     LLM,
+    FileBase,
     FileInput,
     FileWithBase64,
     FileWithId,
-    FileWithUrl,
     InputItem,
     LLMConfig,
-    ProviderConfig,
     QueryResult,
     QueryResultCost,
     QueryResultMetadata,
-    RawInputItem,
+    RawInput,
+    RawResponse,
     TextInput,
     ToolBody,
     ToolCall,
@@ -37,36 +35,25 @@ from model_library.exceptions import (
     BadInputError,
     MaxOutputTokensExceededError,
     ModelNoOutputError,
+    NoMatchingToolCallError,
     RateLimitException,
 )
 from model_library.providers.openai import OpenAIModel
 from model_library.register_models import register_provider
 from model_library.utils import create_openai_client_with_defaults
 
-Chat = AsyncChat | SyncChat
-
-
-class XAIConfig(ProviderConfig):
-    sync_client: bool = False
-
 
 @register_provider("grok")
 class XAIModel(LLM):
-    provider_config = XAIConfig()
-
-    _client: AsyncClient | Client | None = None
+    _client: AsyncClient | None = None
 
     @override
-    def get_client(self) -> AsyncClient | Client:
-        if self._client:
-            return self._client
-
-        ClientClass = Client if self.provider_config.sync_client else AsyncClient
-        self._client = ClientClass(
-            api_key=model_library_settings.XAI_API_KEY,
-        )
-
-        return self._client
+    def get_client(self) -> AsyncClient:
+        if not XAIModel._client:
+            XAIModel._client = AsyncClient(
+                api_key=model_library_settings.XAI_API_KEY,
+            )
+        return XAIModel._client
 
     @override
     def __init__(
@@ -98,54 +85,68 @@ class XAIModel(LLM):
             )
         )
 
+    async def get_tool_call_ids(self, input: Sequence[InputItem]) -> list[str]:
+        raw_responses = [x for x in input if isinstance(x, RawResponse)]
+        tool_call_ids: list[str] = []
+
+        calls = [
+            y
+            for x in raw_responses
+            if isinstance(x.response, Response) and x.response.tool_calls
+            for y in x.response.tool_calls
+        ]
+        tool_call_ids.extend([x.id for x in calls if x.id])
+        return tool_call_ids
+
     @override
     async def parse_input(
         self,
         input: Sequence[InputItem],
         **kwargs: Any,
-    ) -> None:
-        chat: Chat = kwargs["chat"]
+    ) -> list[Message]:
+        new_input: list[Message] = []
+
         content_user: list[Any] = []
+
+        def flush_content_user():
+            if content_user:
+                new_input.append(user(*content_user))
+                content_user.clear()
+
+        tool_call_ids = await self.get_tool_call_ids(input)
+
         for item in input:
+            if isinstance(item, TextInput):
+                content_user.append(item.text)
+                continue
+
+            if isinstance(item, FileBase):
+                match item.type:
+                    case "image":
+                        parsed = await self.parse_image(item)
+                    case "file":
+                        parsed = await self.parse_file(item)
+                content_user.append(parsed)
+                continue
+
+            # non content user item
+            flush_content_user()
+
             match item:
-                case TextInput():
-                    content_user.append(item.text)
-                case FileWithBase64() | FileWithUrl() | FileWithId():
-                    match item.type:
-                        case "image":
-                            content_user.append(await self.parse_image(item))
-                        case "file":
-                            content_user.append(await self.parse_file(item))
-                case _:
-                    if content_user:
-                        chat.append(user(*content_user))
-                        content_user = []
-                    match item:
-                        case ToolResult():
-                            if not (
-                                isinstance(x, Response)
-                                and x.finish_reason == "REASON_TOOL_CALLS"
-                                and x.tool_calls
-                                and any(
-                                    tool
-                                    for tool in x.tool_calls
-                                    if tool.id == item.tool_call.id
-                                )
-                                for x in chat.messages
-                            ):
-                                raise Exception(
-                                    "Tool call result provided with no matching tool call"
-                                )
-                            chat.append(tool_result(item.result))
-                        case dict():  # RawInputItem
-                            item = cast(RawInputItem, item)
-                            chat.append(item)  # pyright: ignore[reportArgumentType]
-                        case _:  # RawResponse
-                            item = cast(Response, item)
-                            chat.append(item)
-
-        if content_user:
-            chat.append(user(*content_user))
+                case ToolResult():
+                    if item.tool_call.id not in tool_call_ids:
+                        raise NoMatchingToolCallError()
+
+                    new_input.append(tool_result(item.result))
+                case RawResponse():
+                    new_input.append(item.response)
+                case RawInput():
+                    new_input.append(item.input)
+
+        # in case content user item is the last item
+        flush_content_user()
+
+        return new_input
 
     @override
     async def parse_image(
@@ -200,43 +201,10 @@ class XAIModel(LLM):
     ) -> FileWithId:
         raise NotImplementedError()
 
-    def fetch_response_sync(
-        self,
-        chat: SyncChat,
-    ) -> Response | None:
-        latest_response = None
-        for response, _ in chat.stream():
-            latest_response = response
-
-        return latest_response
-
-    async def fetch_response_async(
-        self,
-        chat: AsyncChat,
-    ) -> Response | None:
-        latest_response = None
-        async for response, _ in chat.stream():
-            latest_response = response
-
-        return latest_response
-
     @override
-    async def _query_impl(
-        self,
-        input: Sequence[InputItem],
-        *,
-        tools: list[ToolDefinition],
-        query_logger: logging.Logger,
-        **kwargs: object,
-    ) -> QueryResult:
-        if self.reasoning_effort:
-            kwargs["reasoning_effort"] = self.reasoning_effort
-
-        if self.delegate:
-            return await self.delegate_query(
-                input, tools=tools, query_logger=query_logger, **kwargs
-            )
-
+    async def build_body(
+        self, input: Sequence[InputItem], *, tools: list[ToolDefinition], **kwargs: Any
+    ) -> dict[str, Any]:
         messages: Sequence[Message] = []
         if "system_prompt" in kwargs:
             messages.append(system(str(kwargs.pop("system_prompt"))))
@@ -254,20 +222,43 @@ class XAIModel(LLM):
             if self.top_p is not None:
                 body["top_p"] = self.top_p
 
+        if self.reasoning_effort:
+            body["reasoning_effort"] = self.reasoning_effort
+
         body.update(kwargs)
 
+        # use Chat object to parse raw Response and other objects into the correct formats
+        # see xai's chat.py `class BaseChat` -> `def append`
+        chat: Chat = self.get_client().chat.create("dummy")
+        parsed_input = await self.parse_input(input)
+        for message in parsed_input:
+            chat.append(message)
+        body["messages"].extend(chat.messages)
+
+        return body
+
+    @override
+    async def _query_impl(
+        self,
+        input: Sequence[InputItem],
+        *,
+        tools: list[ToolDefinition],
+        query_logger: logging.Logger,
+        **kwargs: object,
+    ) -> QueryResult:
+        if self.delegate:
+            return await self.delegate_query(
+                input, tools=tools, query_logger=query_logger, **kwargs
+            )
+
+        body = await self.build_body(input, tools=tools, **kwargs)
+
         try:
-            chat: Chat = self.get_client().chat.create(**body)  # pyright: ignore[reportAny]
-            await self.parse_input(input, chat=chat)
-
-            # Allows users to dynamically swap to a sync client if getting grpc errors
-            # Run in a separate thread so we are playing fair with other async processes
-            if self.provider_config.sync_client:
-                latest_response = await asyncio.to_thread(
-                    self.fetch_response_sync, cast(SyncChat, chat)
-                )
-            else:
-                latest_response = await self.fetch_response_async(cast(AsyncChat, chat))
+            chat: Chat = self.get_client().chat.create(**body)
+
+            latest_response: Response | None = None
+            async for response, _ in chat.stream():
+                latest_response = response
 
             if not latest_response:
                 raise ModelNoOutputError("Model failed to produce a response")
@@ -307,8 +298,25 @@ class XAIModel(LLM):
                 cache_read_tokens=latest_response.usage.cached_prompt_text_tokens,
             ),
             tool_calls=tool_calls,
-            history=[*input, latest_response],
+            history=[*input, RawResponse(response=latest_response)],
+        )
+
+    @override
+    async def count_tokens(
+        self,
+        input: Sequence[InputItem],
+        *,
+        history: Sequence[InputItem] = [],
+        tools: list[ToolDefinition] = [],
+        **kwargs: object,
+    ) -> int:
+        string_input = await self.stringify_input(input, history=history, tools=tools)
+        self.logger.debug(string_input)
+
+        tokens = await self.get_client().tokenize.tokenize_text(
+            string_input, self.model_name
         )
+        return len(tokens)
 
     @override
     async def _calculate_cost(

model_library/register_models.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from typing import Any, Callable, Type, TypeVar, cast, get_type_hints
 
 import yaml
-from pydantic import create_model, model_validator
+from pydantic import ConfigDict, create_model, model_validator
 from pydantic.fields import Field
 from pydantic.main import BaseModel
 
@@ -173,6 +173,8 @@ class DefaultParameters(BaseModel):
 
 
 class RawModelConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+
     company: str
     label: str
     description: str | None = None
@@ -275,7 +277,7 @@ def _register_models() -> ModelRegistry:
 
                     # create model config object
                     raw_model_obj: RawModelConfig = RawModelConfig.model_validate(
-                        current_model_config, strict=True, extra="forbid"
+                        current_model_config
                     )
 
                     provider_endpoint = (

model_library/utils.py

@@ -1,6 +1,5 @@
 import logging
 from collections.abc import Mapping, Sequence
-from typing import Any
 
 import httpx
 from openai import AsyncOpenAI
@@ -90,37 +89,3 @@ def get_context_window_for_model(model_name: str, default: int = 128_000) -> int
             f"using default context length of {default}"
         )
         return default
-
-
-def normalize_tool_result(result: Any) -> str:
-    """Normalize tool result to non-empty string for API compatibility.
-
-    Empty results (None, empty dict/list, whitespace-only strings) are
-    converted to a single space to satisfy API requirements.
-
-    Args:
-        result: Tool result value (any type)
-
-    Returns:
-        Non-empty string representation of the result
-    """
-    if result is None or (isinstance(result, (dict, list)) and not result):
-        return " "
-    result_str = str(result)  # pyright: ignore[reportUnknownArgumentType]
-    return result_str.strip() or " "
-
-
-def filter_empty_text_blocks(content: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    """Filter out empty text blocks from content list.
-
-    Args:
-        content: List of content blocks (dicts with 'type' and potentially 'text' keys)
-
-    Returns:
-        Filtered list with empty text blocks removed
-    """
-    return [
-        block
-        for block in content
-        if block.get("type") != "text" or block.get("text", "").strip()
-    ]

{model_library-0.1.6.dist-info → model_library-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-library
-Version: 0.1.6
+Version: 0.1.7
 Summary: Model Library for vals.ai 
 Author-email: "Vals AI, Inc." <contact@vals.ai>
 License: MIT
@@ -13,13 +13,13 @@ Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: rich
 Requires-Dist: backoff<3.0,>=2.2.1
 Requires-Dist: redis<7.0,>=6.2.0
-Requires-Dist: tiktoken==0.11.0
+Requires-Dist: tiktoken>=0.12.0
 Requires-Dist: pillow
 Requires-Dist: openai<3.0,>=2.0
 Requires-Dist: anthropic<1.0,>=0.57.1
 Requires-Dist: mistralai<2.0,>=1.9.10
 Requires-Dist: xai-sdk<2.0,>=1.0.0
-Requires-Dist: ai21<5.0,>=4.0.3
+Requires-Dist: ai21<5.0,>=4.3.0
 Requires-Dist: boto3<2.0,>=1.38.27
 Requires-Dist: google-genai[aiohttp]>=1.51.0
 Requires-Dist: google-cloud-storage>=1.26.0

{model_library-0.1.6.dist-info → model_library-0.1.7.dist-info}/RECORD

@@ -1,20 +1,20 @@
 model_library/__init__.py,sha256=AKc_15aklOf-LbcS9z1Xer_moRWNpG6Dh3kqvSQ0nOI,714
-model_library/exceptions.py,sha256=ZHMr6lloXZz4V4Wy1UP8zc1CdUHx6-IS9_rOi6oN45s,8680
+model_library/exceptions.py,sha256=4TJ1aDkpPV-gv3gLIO7pi5ORRBG2hPXSBAvOXS6I_Wg,9027
 model_library/file_utils.py,sha256=FAZRRtDT8c4Rjfoj64Te3knEHggXAAfRRuS8WLCsSe8,3682
-model_library/logging.py,sha256=McyaPHUk7RkB38-LrfnudrrU1B62ta8wAbbIBwLRmj0,853
+model_library/logging.py,sha256=rZrrVQlEmyZzvKx6nIOR8bKHl49wQIIW5c36Zqcigm4,888
 model_library/model_utils.py,sha256=l8oCltGeimMGtnne_3Q1EguVtzCj61UMsLsma-1czwg,753
 model_library/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-model_library/register_models.py,sha256=3FeFrcS2qRpAhj9ahXNuZ6jcH5UEks3I_PaT6rPvKgs,13653
+model_library/register_models.py,sha256=-CtFggD296wqO2nd6yC5GIWBoi0YD0kDyG8bwxBa1ec,13696
 model_library/registry_utils.py,sha256=BVauHcP02Et2maLxowNBbdpGd32cnLz1_zSjDLVJjp0,8843
 model_library/settings.py,sha256=QyeUqzWBpexFi014L_mZkoXP49no3SAQNJRObATXrL8,873
-model_library/utils.py,sha256=T91ACGTc-KtksVyMFspt-vJtR5I-xcO3nVfH6SltmMU,3988
+model_library/utils.py,sha256=dcww26Ph0PSihEzSN4gZ48XXxN3Bt5oWveVYX6UACyI,2922
 model_library/base/__init__.py,sha256=TtxCXGUtkEqWZNMMofLPuC4orN7Ja2hemtbtHitt_UA,266
-model_library/base/base.py,sha256=mvubt5VC1eM8cuLw_RHP04hTgNEcULzIBiJcHqKF--c,14289
+model_library/base/base.py,sha256=nrdPS92-vFkDVJJvNVOQ8e6U1zCyZxOtFg18ZgM4o0s,17467
 model_library/base/batch.py,sha256=-jd6L0ECc5pkj73zoX2ZYcv_9iQdqxEi1kEilwaXWSA,2895
-model_library/base/delegate_only.py,sha256=YJUvP9k8x2kBsI-0ACNmx1Jx77zdZSBWCMSpx0LAyXE,2213
-model_library/base/input.py,sha256=Nhg8Ril1kFau1DnE8u102JC1l-vxNd-v9e3SjovR-Do,1876
-model_library/base/output.py,sha256=Ak6CJRYqtjYILsSWkfE70fSK3yvP7v_n5NYfysMaIL4,7464
-model_library/base/utils.py,sha256=YGQLPyQgCbfHNBxyTxCvpZNZ-ctEji258IdfMiXUJXs,1962
+model_library/base/delegate_only.py,sha256=5v2twEuQl1jF34M8iFcaZQlk0_uBLew4B46TkHUEssw,2441
+model_library/base/input.py,sha256=JrnvBZ_xLcEmaMjnOfUS6GFV0QWtCGpJq0RQQL2YBG8,1934
+model_library/base/output.py,sha256=jwd3rfRFUcqm8q-O5H684ToNZPvcD_obtf6ugmhzUus,7613
+model_library/base/utils.py,sha256=eiMTiFFXHTb44Nnz3fjxf9YQzJpdfNI7tprSi9COPu0,2268
 model_library/config/README.md,sha256=i8_wHnlI6uHIqWN9fYBkDCglZM2p5ZMVD3SLlxiwUVk,4274
 model_library/config/ai21labs_models.yaml,sha256=ZWHhk1cep2GQIYHqkTS_0152mF3oZg2tSzMPmvfMRSI,2478
 model_library/config/alibaba_models.yaml,sha256=-RLWOwh3ZaCQqjaZ-4Zw0BJNVE6JVHJ8Ggm9gQJZ6QI,2082
@@ -36,29 +36,29 @@ model_library/config/together_models.yaml,sha256=BeqRJonYzPvWwoLfkyH0RMRKBYUrCSE
 model_library/config/xai_models.yaml,sha256=2KRNNQy3kV-4xxSfhj7Uhp9TZF-S5qPlM8Ef-04zv8Y,7985
 model_library/config/zai_models.yaml,sha256=lcYMh2FCrLWkKqdCnarRlwDoL3SbutRBNAiMPBUYQiw,1894
 model_library/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-model_library/providers/ai21labs.py,sha256=sqmu9R7owZZQLxOkNV9dhSeZVAlTMDCNFVdxJyJo6UA,5885
+model_library/providers/ai21labs.py,sha256=NXoEu0NFwxLVJIN751tEuKLFYk4ECSDnWMik35pL5Y4,6277
 model_library/providers/alibaba.py,sha256=k6LZErV_l9oTFTdKTwyw1SXD509Rl3AqFbN8umCryEE,2941
-model_library/providers/amazon.py,sha256=U0tH5mw8dAMDg06BtnVoR-RxYlcAL1-7ZR737sR8xgU,14068
-model_library/providers/anthropic.py,sha256=G94hFkRFTWutq9XYd3321KkPrxVHnR6fs_h2AdkVqx4,23197
+model_library/providers/amazon.py,sha256=F8anL4pmeadA0P9luUZSSnEcke90WH7mCa3afJ2akgM,14156
+model_library/providers/anthropic.py,sha256=xbSZx6R54UhNQij-vwFKPCIQJoTzxLblJMIxdqGcsqs,23094
 model_library/providers/azure.py,sha256=brQNCED-zHvYjL5K5hdjFBNso6hJZg0HTHNnAgJPPG0,1408
 model_library/providers/cohere.py,sha256=lCBm1PP1l_UOa1pKFMIZM3C0wCv3QWB6UP0-jvjkFa4,1066
 model_library/providers/deepseek.py,sha256=7T4lxDiV5wmWUK7TAKwr332_T6uyXNCOiirZOCCETL0,1159
 model_library/providers/fireworks.py,sha256=w-5mOF5oNzqx_0ijCoTm1lSn2ZHwhp6fURKhV3LEqIc,2309
 model_library/providers/inception.py,sha256=Nrky53iujIM9spAWoNRtoJg2inFiL0li6E75vT3b6V8,1107
 model_library/providers/kimi.py,sha256=zzvcKpZLsM1xPebpLeMxNKTt_FRiLN1rFWrIly7wfXA,1092
-model_library/providers/minimax.py,sha256=YRtJW2wgiu6KXEBScYETeVMNTfhPvpjL2J-oo0wE_BI,1057
-model_library/providers/mistral.py,sha256=r0PY30kHY-guaSzIEahdp2I45meJzo71Ql97NfkPv-8,9798
-model_library/providers/openai.py,sha256=MMm6K4iewhSpPzEeRhrPRYf_txrpklCrefNHiUly8S8,33665
+model_library/providers/minimax.py,sha256=ckVyoXdtVxGT3aU-AknBQCa7_mOckNMfOXSgwbrJNIY,1610
+model_library/providers/mistral.py,sha256=kmp74jEETMnB8fQ5VNfNVkksIrPMGJwJeXJDnTVwKa8,10117
+model_library/providers/openai.py,sha256=gMykcN5eHFHqCrK9y5twc18h7XSMI68mB9YGYnpt93A,34522
 model_library/providers/perplexity.py,sha256=eIzzkaZ4ZMlRKFVI9bnwyo91iJkh7aEmJ-0_4OKeAWc,1083
 model_library/providers/together.py,sha256=7Y4QLnX8c_fyXUud-W_C1gidmROQainTgODBwbvFyXQ,2033
-model_library/providers/vals.py,sha256=mKaItg_g9RJeaIDhoBu7ksTe42P0MRYFI4X1cla8YC0,9883
-model_library/providers/xai.py,sha256=toSqWBHUaHE000aMdOayAW3-_ZmDUotWEpZ4-X33LuY,10918
+model_library/providers/vals.py,sha256=dbaL8toYTssm8qVlrgqzqwCeeTV9h-xfA37uBuoWtfg,9894
+model_library/providers/xai.py,sha256=eA3CsiOWIF_vKxNZC95INUKOazz54R6vpVrGM8VR1tY,10719
 model_library/providers/zai.py,sha256=O_GM6KlJ0fM2wYoxO9xrCWfnpYH7IpoKEzjiD4jB8Kc,1050
 model_library/providers/google/__init__.py,sha256=ypuLVL_QJEQ7C3S47FhC9y4wyawYOdGikAViJmACI0U,115
-model_library/providers/google/batch.py,sha256=4TE90Uo1adi54dVtGcGyUAxw11YExJq-Y4KmkQ-cyHA,9978
-model_library/providers/google/google.py,sha256=xmiktN-Z9W1fC1jHUT_m6x5fTpI6-mWpKvbMGg9kgXE,16787
-model_library-0.1.6.dist-info/licenses/LICENSE,sha256=x6mf4o7U_wHaaqcfxoU-0R6uYJLbqL_TNuoULP3asaA,1070
-model_library-0.1.6.dist-info/METADATA,sha256=sNWBOgDqydFI184UERputqhulBz0olrbye-fO7owrCE,6989
-model_library-0.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-model_library-0.1.6.dist-info/top_level.txt,sha256=HtQYxA_7RP8UT35I6VcUw20L6edI0Zf2t5Ys1uDGVjs,14
-model_library-0.1.6.dist-info/RECORD,,
+model_library/providers/google/batch.py,sha256=ycK2arf00lhZQNgyM5Yd01LAScul6rvVv--dUcRWWSA,9977
+model_library/providers/google/google.py,sha256=txQaet1HobyjYd3dp9Mgonx3x5ln3LMuW10Qsyum3B4,17809
+model_library-0.1.7.dist-info/licenses/LICENSE,sha256=x6mf4o7U_wHaaqcfxoU-0R6uYJLbqL_TNuoULP3asaA,1070
+model_library-0.1.7.dist-info/METADATA,sha256=436seE0dN2884VkHzGHyp36aARqBxqKMKRccmKYG3_E,6989
+model_library-0.1.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+model_library-0.1.7.dist-info/top_level.txt,sha256=HtQYxA_7RP8UT35I6VcUw20L6edI0Zf2t5Ys1uDGVjs,14
+model_library-0.1.7.dist-info/RECORD,,