🐛 FIX: Types

Author: arre-ankit

langbase/__init__.py

@@ -45,6 +45,7 @@
 from .types import (
     ChoiceGenerate,
     Message,
+    MessageContentItem,
     PipeBaseOptions,
     PipeBaseResponse,
     PipeCreateOptions,

langbase/primitives/agent.py

@@ -2,10 +2,11 @@
 Agent API client for the Langbase SDK.
 """
 
-from typing import Any, Dict, List, Optional, Union, overload
+from typing import Any, Dict, List, Literal, Optional, Union, overload
 
 from langbase.constants import AGENT_RUN_ENDPOINT
 from langbase.request import Request
+from langbase.types import McpServerSchema, Message, ToolChoice, Tools
 from langbase.utils import clean_null_values
 
 
@@ -17,7 +18,7 @@ def __init__(self, parent):
     @overload
     def run(
         self,
-        input: Union[str, List[Dict[str, Any]]],
+        input: Union[str, List[Message]],
         model: str,
         api_key: str,
         instructions: Optional[str] = None,
@@ -27,14 +28,14 @@ def run(
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         stop: Optional[List[str]] = None,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Tools]] = None,
+        tool_choice: Optional[Union[Literal["auto", "required"], ToolChoice]] = None,
         parallel_tool_calls: Optional[bool] = None,
         reasoning_effort: Optional[str] = None,
         max_completion_tokens: Optional[int] = None,
         response_format: Optional[Dict[str, Any]] = None,
         custom_model_params: Optional[Dict[str, Any]] = None,
-        mcp_servers: Optional[List[Dict[str, Any]]] = None,
+        mcp_servers: List[McpServerSchema] = None,
         *,
         stream: bool = True,
     ) -> Any:
@@ -44,7 +45,7 @@ def run(
     @overload
     def run(
         self,
-        input: Union[str, List[Dict[str, Any]]],
+        input: Union[str, List[Message]],
         model: str,
         api_key: str,
         instructions: Optional[str] = None,
@@ -54,22 +55,22 @@ def run(
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         stop: Optional[List[str]] = None,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Tools]] = None,
+        tool_choice: Optional[Union[Literal["auto", "required"], ToolChoice]] = None,
         parallel_tool_calls: Optional[bool] = None,
         reasoning_effort: Optional[str] = None,
         max_completion_tokens: Optional[int] = None,
         response_format: Optional[Dict[str, Any]] = None,
         custom_model_params: Optional[Dict[str, Any]] = None,
-        mcp_servers: Optional[List[Dict[str, Any]]] = None,
+        mcp_servers: List[McpServerSchema] = None,
         stream: bool = False,
     ) -> Dict[str, Any]:
         """Non-stream overload - returns dict response when stream=False"""
         ...
 
     def run(
         self,
-        input: Union[str, List[Dict[str, Any]]],
+        input: Union[str, List[Message]],
         model: str,
         api_key: str,
         instructions: Optional[str] = None,
@@ -79,14 +80,14 @@ def run(
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         stop: Optional[List[str]] = None,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        tools: Optional[List[Tools]] = None,
+        tool_choice: Optional[Union[Literal["auto", "required"], ToolChoice]] = None,
         parallel_tool_calls: Optional[bool] = None,
         reasoning_effort: Optional[str] = None,
         max_completion_tokens: Optional[int] = None,
         response_format: Optional[Dict[str, Any]] = None,
         custom_model_params: Optional[Dict[str, Any]] = None,
-        mcp_servers: Optional[List[Dict[str, Any]]] = None,
+        mcp_servers: List[McpServerSchema] = None,
         stream: bool = False,
     ) -> Union[Dict[str, Any], Any]:
         """

langbase/primitives/pipes.py

@@ -7,6 +7,7 @@
 from langbase.constants import PIPE_DETAIL_ENDPOINT, PIPE_RUN_ENDPOINT, PIPES_ENDPOINT
 from langbase.request import Request
 from langbase.types import (
+    Message,
     PipeCreateResponse,
     PipeListResponse,
     PipeUpdateResponse,
@@ -65,9 +66,9 @@ def update(self, name: str, **kwargs) -> PipeUpdateResponse:
 
     def run(
         self,
-        name: Optional[str] = None,
+        name: str = None,
         api_key: Optional[str] = None,
-        messages: Optional[List[Dict[str, Any]]] = None,
+        messages: Optional[List[Message]] = None,
         variables: Optional[List[Dict[str, str]]] = None,
         thread_id: Optional[str] = None,
         raw_response: Optional[bool] = None,

langbase/types.py

@@ -79,11 +79,18 @@ class ToolChoice(TypedDict):
     function: Dict[str, str]
 
 
+class MessageContentItem(TypedDict, total=False):
+    type: str
+    text: Optional[str]
+    image_url: Optional[Dict[str, str]]
+    cache_control: Optional[Dict[str, str]]
+
+
 class Message(TypedDict, total=False):
     """Basic message structure."""
 
     role: Role
-    content: Optional[str]
+    content: Union[str, List[MessageContentItem], None]
     name: Optional[str]
     tool_call_id: Optional[str]
     tool_calls: Optional[List[ToolCall]]