openrouter-provider 0.0.5__py3-none-any.whl → 1.0.10__py3-none-any.whl

openrouter/openrouter_provider.py

@@ -0,0 +1,279 @@
+from __future__ import annotations
+import json
+import os
+from dataclasses import dataclass, asdict
+from typing import List, Optional, Literal, Iterator, AsyncIterator
+
+from dotenv import load_dotenv
+from openai import OpenAI, AsyncOpenAI
+from openai.types.chat import ChatCompletionChunk
+from pydantic import BaseModel, ValidationError
+
+from openrouter.message import Message, Role, ToolCall
+from openrouter.tool import tool_model
+from openrouter.llms import LLMModel
+
+
+
+
+@dataclass
+class ProviderConfig:
+    order: Optional[List[str]] = None
+    allow_fallbacks: bool = None
+    require_parameters: bool = None
+    data_collection: Literal["allow", "deny"] = None
+    only: Optional[List[str]] = None
+    ignore: Optional[List[str]] = None
+    quantizations: Optional[List[str]] = None
+    sort: Optional[Literal["price", "throughput"]] = None
+    max_price: Optional[dict] = None
+
+    def to_dict(self) -> dict:
+        return {k: v for k, v in asdict(self).items() if v is not None}
+
+
+class OpenRouterProvider:
+    def __init__(self) -> None:
+        load_dotenv()
+        api_key = os.getenv("OPENROUTER_API_KEY")
+        if not api_key:
+            raise ValueError("OPENROUTER_API_KEY is not set in environment variables.")
+        self.client = OpenAI(
+            base_url="https://openrouter.ai/api/v1",
+            api_key=api_key,
+        )
+        self.async_client = AsyncOpenAI(
+            base_url="https://openrouter.ai/api/v1",
+            api_key=api_key,
+        )
+
+    def make_prompt(
+        self,
+        system_prompt: Message,
+        querys: list[Message]
+    ) -> list[dict]:
+        messages = [{"role": "system", "content": system_prompt.text}]
+
+        for query in querys:
+            if query.role == Role.user:
+                if query.images is None:
+                    messages.append({"role": "user", "content": query.text})
+                else:
+                    content = [{"type": "text", "text": query.text}]
+                    for img in query.images[:50]:
+                        content.append(
+                            {"type": "image_url",
+                             "image_url": {"url": f"data:image/jpeg;base64,{img}"}})
+                    messages.append({"role": "user", "content": content})
+
+            elif query.role == Role.ai or query.role == Role.tool:
+                assistant_msg = {"role": "assistant"}
+                assistant_msg["content"] = query.text or None
+
+                if query.tool_calls:
+                    assistant_msg["tool_calls"] = [
+                        {
+                            "id": str(t.id),
+                            "type": "function",
+                            "function": {
+                                "name": t.name,
+                                "arguments": t.arguments
+                            }
+                        }
+                        for t in query.tool_calls
+                    ]
+                messages.append(assistant_msg)
+
+            for t in query.tool_calls:
+                messages.append({
+                    "role": "tool",
+                    "tool_call_id": str(t.id),
+                    "content": str(t.result)
+                })
+            
+        return messages
+
+    def invoke(
+        self,
+        model: LLMModel,
+        system_prompt: Message,
+        querys: list[Message],
+        tools: list[tool_model] = None,
+        provider: ProviderConfig = None,
+        temperature: float = 0.3
+    ) -> Message:
+        tools = tools or []
+        messages = self.make_prompt(system_prompt, querys)
+
+        tool_defs = [tool.tool_definition for tool in tools] if tools else None
+        provider_dict = provider.to_dict() if provider else None
+        
+        response = self.client.chat.completions.create(
+            model=model.name,
+            temperature=temperature,
+            messages=messages,
+            tools=tool_defs,
+            extra_body={"provider": provider_dict},
+        )
+
+        reply = Message(text=response.choices[0].message.content, role=Role.ai, raw_response=response)
+
+        if response.choices[0].message.tool_calls:
+            reply.role = Role.tool
+            for tool in response.choices[0].message.tool_calls:
+                reply.tool_calls.append(ToolCall(id=tool.id, name=tool.function.name, arguments=tool.function.arguments))
+        return reply
+    
+    def invoke_stream(
+        self,
+        model: LLMModel,
+        system_prompt: Message,
+        querys: list[Message],
+        tools: list[tool_model] = None,
+        provider: ProviderConfig = None,
+        temperature: float = 0.3
+    ) -> Iterator[ChatCompletionChunk]:
+        tools = tools or []
+        messages = self.make_prompt(system_prompt, querys)
+
+        tool_defs = [tool.tool_definition for tool in tools] if tools else None
+        provider_dict = provider.to_dict() if provider else None
+
+        response = self.client.chat.completions.create(
+            model=model.name,
+            temperature=temperature,
+            messages=messages,
+            tools=tool_defs,
+            extra_body={"provider": provider_dict},
+            stream=True
+        )
+        
+        return response
+
+    async def async_invoke(
+        self,
+        model: LLMModel,
+        system_prompt: Message,
+        querys: list[Message],
+        tools: list[tool_model] = None,
+        provider: ProviderConfig = None,
+        temperature: float = 0.3
+    ) -> Message:
+        tools = tools or []
+        messages = self.make_prompt(system_prompt, querys)
+
+        tool_defs = [tool.tool_definition for tool in tools] if tools else None
+        provider_dict = provider.to_dict() if provider else None
+
+        response = await self.async_client.chat.completions.create(
+            model=model.name,
+            temperature=temperature,
+            messages=messages,
+            tools=tool_defs,
+            extra_body={"provider": provider_dict}
+        )
+
+        reply = Message(text=response.choices[0].message.content, role=Role.ai, raw_response=response)
+
+        if response.choices[0].message.tool_calls:
+            reply.role = Role.tool
+            for tool in response.choices[0].message.tool_calls:
+                reply.tool_calls.append(ToolCall(id=tool.id, name=tool.function.name, arguments=tool.function.arguments))
+        return reply
+        
+    async def async_invoke_stream(
+        self,
+        model: LLMModel,
+        system_prompt: Message,
+        querys: list[Message],
+        tools: list[tool_model] = None,
+        provider: ProviderConfig = None,
+        temperature: float = 0.3
+    ) -> AsyncIterator[ChatCompletionChunk]:
+        tools = tools or []
+        messages = self.make_prompt(system_prompt, querys)
+
+        tool_defs = [tool.tool_definition for tool in tools] if tools else None
+        provider_dict = provider.to_dict() if provider else None
+
+        response = await self.async_client.chat.completions.create(
+            model=model.name,
+            temperature=temperature,
+            messages=messages,
+            tools=tool_defs,
+            extra_body={"provider": provider_dict},
+            stream=True
+        )
+
+        async for chunk in response:
+            yield chunk
+        
+    def structured_output(
+        self,
+        model: LLMModel,
+        system_prompt: Message,
+        querys: list[Message],
+        provider: ProviderConfig = None,
+        json_schema: BaseModel = None,
+        temperature: float = 0.3
+    ) -> BaseModel:
+        messages = self.make_prompt(system_prompt, querys)
+        provider_dict = provider.to_dict() if provider else None
+        
+        schema = json_schema.model_json_schema()
+        
+        def add_additional_properties_false(obj):
+            if isinstance(obj, dict):
+                if "properties" in obj:
+                    obj["additionalProperties"] = False
+                for value in obj.values():
+                    add_additional_properties_false(value)
+            elif isinstance(obj, list):
+                for item in obj:
+                    add_additional_properties_false(item)
+        
+        def ensure_required_properties(obj):
+            if isinstance(obj, dict):
+                properties = obj.get("properties")
+                if isinstance(properties, dict):
+                    keys = list(properties.keys())
+                    existing_required = obj.get("required")
+                    if isinstance(existing_required, list):
+                        required_set = set(existing_required)
+                    else:
+                        required_set = set()
+                    required_set.update(keys)
+                    obj["required"] = list(required_set)
+                for value in obj.values():
+                    ensure_required_properties(value)
+            elif isinstance(obj, list):
+                for item in obj:
+                    ensure_required_properties(item)
+
+        add_additional_properties_false(schema)
+        ensure_required_properties(schema)
+        
+        response = self.client.chat.completions.create(
+            model=model.name,
+            temperature=temperature,
+            messages=messages,
+            response_format={"type": "json_schema", "json_schema": {"name": json_schema.__name__, "schema": schema}},
+            extra_body={"provider": provider_dict},
+        )
+
+        content = response.choices[0].message.content
+
+        try:
+            return json_schema.model_validate_json(content)
+        except ValidationError:
+            formatted_content = content
+            try:
+                parsed = json.loads(content)
+                formatted_content = json.dumps(parsed, indent=2, ensure_ascii=False)
+            except json.JSONDecodeError:
+                pass
+            print("structured_output validation failed, response content:")
+            print(formatted_content)
+            raise
+
+    

OpenRouterProvider/Tool.py → openrouter/tool.py

@@ -1,10 +1,11 @@
+from __future__ import annotations
 import inspect
 from functools import wraps
-from typing import get_type_hints, get_origin, get_args
+from typing import get_type_hints, get_origin, get_args, Callable, Any, Dict
 
 
 class tool_model:
-    def __init__(self, func):
+    def __init__(self, func: Callable) -> None:
         self.func = func
         self.name = func.__name__
 
@@ -12,22 +13,20 @@ class tool_model:
         type_hints = get_type_hints(func)
 
         type_map = {
-            int:   "integer",
-            str:   "string",
+            int: "integer",
+            str: "string",
             float: "number",
-            bool:  "boolean",
+            bool: "boolean",
         }
 
-        properties = {}
-        required = []
+        properties: Dict[str, Dict[str, Any]] = {}
+        required: list[str] = []
 
-        # Build JSON schema properties
         for name, param in sig.parameters.items():
             anno = type_hints.get(name, None)
             origin = get_origin(anno)
 
             if origin is list:
-                # Handle list element types
                 (elem_type,) = get_args(anno)
                 json_item_type = type_map.get(elem_type, "string")
                 schema = {
@@ -36,7 +35,6 @@ class tool_model:
                     "description": name,
                 }
             else:
-                # Primitive types
                 json_type = type_map.get(anno, "string")
                 schema = {
                     "type": json_type,
@@ -47,7 +45,6 @@ class tool_model:
             if param.default is inspect._empty:
                 required.append(name)
 
-        # Attach tool definition metadata
         self.tool_definition = {
             "type": "function",
             "function": {
@@ -65,7 +62,7 @@ class tool_model:
 
         wraps(func)(self)
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
         return self.func(*args, **kwargs)