npcsh 0.3.31__py3-none-any.whl → 0.3.32__py3-none-any.whl

npcsh/stream.py

@@ -6,397 +6,41 @@
 ########
 ########
 
-from npcsh.npc_sysenv import get_system_message
 from typing import Any, Dict, Generator, List
 import os
-import anthropic
-from openai import OpenAI
-from google import genai
-
-from google.generativeai import types
-import google.generativeai as genai
 import base64
 import json
+import requests
+from PIL import Image
+from typing import Any, Dict, Generator, List, Union
 
+from pydantic import BaseModel
+from npcsh.npc_sysenv import (
+    get_system_message,
+    compress_image,
+    available_chat_models,
+    available_reasoning_models,
+)
 
-def get_anthropic_stream(
-    messages,
-    model: str,
-    npc: Any = None,
-    tools: list = None,
-    images: List[Dict[str, str]] = None,
-    api_key: str = None,
-    tool_choice: Dict = None,
-    **kwargs,
-) -> Generator:
-    """
-    Streams responses from Anthropic, supporting images, tools, and yielding raw text chunks.
-
-    Args:
-        messages: List of conversation messages
-        model: Anthropic model to use
-        npc: Optional NPC context
-        tools: Optional list of tools to provide to Claude
-        images: Optional list of images to include
-        api_key: Anthropic API key
-        tool_choice: Optional tool choice configuration
-        **kwargs: Additional arguments for the API call
-    """
-    if api_key is None:
-        api_key = os.environ.get("ANTHROPIC_API_KEY")
-    client = anthropic.Anthropic(api_key=api_key)
-
-    if messages[0]["role"] == "system":
-        system_message = messages[0]["content"]
-        messages = messages[1:]
-    elif npc is not None:
-        system_message = get_system_message(npc)
-    else:
-        system_message = "You are a helpful assistant."
-
-    # Preprocess messages to ensure content is a list of dicts
-    for message in messages:
-        if isinstance(message["content"], str):
-            message["content"] = [{"type": "text", "text": message["content"]}]
-    # Add images if provided
-    if images:
-        for img in images:
-            with open(img["file_path"], "rb") as image_file:
-                img["data"] = base64.b64encode(image_file.read()).decode("utf-8")
-                img["media_type"] = "image/jpeg"
-                messages[-1]["content"].append(
-                    {
-                        "type": "image",
-                        "source": {
-                            "type": "base64",
-                            "media_type": img["media_type"],
-                            "data": img["data"],
-                        },
-                    }
-                )
-
-    # Prepare API call parameters
-    api_params = {
-        "model": model,
-        "messages": messages,
-        "max_tokens": kwargs.get("max_tokens", 8192),
-        "stream": True,
-        "system": system_message,
-    }
-
-    # Add tools if provided
-    if tools:
-        api_params["tools"] = tools
-
-    # Add tool choice if specified
-    if tool_choice:
-        api_params["tool_choice"] = tool_choice
-
-    # Make the API call
-    response = client.messages.create(**api_params)
-
-    for chunk in response:
-        yield chunk
-
-
-def process_anthropic_tool_stream(
-    stream, tool_map: Dict[str, callable], messages: List[Dict] = None
-) -> List[Dict]:
-    """
-    Process the Anthropic tool use stream
-    """
-    tool_results = []
-    current_tool = None
-    current_input = ""
-    context = messages[-1]["content"] if messages else ""
-
-    for chunk in stream:
-        # Look for tool use blocks
-        if (
-            chunk.type == "content_block_start"
-            and getattr(chunk, "content_block", None)
-            and chunk.content_block.type == "tool_use"
-        ):
-            current_tool = {
-                "id": chunk.content_block.id,
-                "name": chunk.content_block.name,
-            }
-            current_input = ""
-
-        # Collect input JSON deltas
-        if chunk.type == "content_block_delta" and hasattr(chunk.delta, "partial_json"):
-            current_input += chunk.delta.partial_json
-
-        # When tool input is complete
-        if chunk.type == "content_block_stop" and current_tool:
-            try:
-                # Parse the complete input
-                tool_input = json.loads(current_input) if current_input.strip() else {}
-
-                # Add context to tool input
-                tool_input["context"] = context
-
-                # Execute the tool
-                tool_func = tool_map.get(current_tool["name"])
-                if tool_func:
-                    result = tool_func(tool_input)
-                    tool_results.append(
-                        {
-                            "tool_name": current_tool["name"],
-                            "tool_input": tool_input,
-                            "tool_result": result,
-                        }
-                    )
-                else:
-                    tool_results.append(
-                        {
-                            "tool_name": current_tool["name"],
-                            "tool_input": tool_input,
-                            "error": f"Tool {current_tool['name']} not found",
-                        }
-                    )
-
-            except Exception as e:
-                tool_results.append(
-                    {
-                        "tool_name": current_tool["name"],
-                        "tool_input": current_input,
-                        "error": str(e),
-                    }
-                )
-
-            # Reset current tool
-            current_tool = None
-            current_input = ""
-
-    return tool_results
-
-
-from typing import List, Dict, Any, Literal
-
-ProviderType = Literal[
-    "openai", "anthropic", "ollama", "gemini", "deepseek", "openai-like"
-]
-
-
-def generate_tool_schema(
-    name: str,
-    description: str,
-    parameters: Dict[str, Any],
-    provider: ProviderType,
-    required: List[str] = None,
-) -> Dict[str, Any]:
-    """
-    Generate provider-specific function/tool schema from common parameters
-
-    Args:
-        name: Name of the function
-        description: Description of what the function does
-        parameters: Dict of parameter names and their properties
-        provider: Which provider to generate schema for
-        required: List of required parameter names
-    """
-    if required is None:
-        required = []
-
-    if provider == "openai":
-        return {
-            "type": "function",
-            "function": {
-                "name": name,
-                "description": description,
-                "parameters": {
-                    "type": "object",
-                    "properties": parameters,
-                    "required": required,
-                    "additionalProperties": False,
-                },
-                "strict": True,
-            },
-        }
-
-    elif provider == "anthropic":
-        return {
-            "name": name,
-            "description": description,
-            "parameters": {
-                "type": "object",
-                "properties": parameters,
-                "required": required,
-            },
-        }
-
-    elif provider == "ollama":
-        return {
-            "type": "function",
-            "function": {
-                "name": name,
-                "description": description,
-                "parameters": {
-                    "type": "object",
-                    "properties": parameters,
-                    "required": required,
-                },
-            },
-        }
-    elif provider == "gemini":
-        # Convert our generic tool schema to a Gemini function declaration
-        function = {
-            "name": name,
-            "description": description,
-            "parameters": {
-                "type": "OBJECT",
-                "properties": {
-                    k: {
-                        "type": v.get("type", "STRING").upper(),
-                        "description": v.get("description", ""),
-                        "enum": v.get("enum", None),
-                    }
-                    for k, v in parameters.items()
-                },
-                "required": required or [],
-            },
-        }
-
-        # Create a Tool object as shown in the example
-        return types.Tool(function_declarations=[function])
-
-    raise ValueError(f"Unknown provider: {provider}")
-
-
-def get_ollama_stream(
-    messages: List[Dict[str, str]],
-    model: str,
-    npc: Any = None,
-    tools: list = None,
-    images: list = None,
-    tool_choice: Dict = None,
-    **kwargs,
-) -> Generator:
-    """Streams responses from Ollama, supporting images and tools."""
-    import ollama
-
-    messages_copy = messages.copy()
-
-    # Handle images if provided
-    if images:
-        messages[-1]["images"] = [image["file_path"] for image in images]
-
-    # Add system message if not present
-    if messages_copy[0]["role"] != "system":
-        if npc is not None:
-            system_message = get_system_message(npc)
-            messages_copy.insert(0, {"role": "system", "content": system_message})
-
-    # Prepare API call parameters
-    api_params = {
-        "model": model,
-        "messages": messages_copy,
-        "stream": True,
-    }
-
-    # Add tools if provided
-    if tools:
-        api_params["tools"] = tools
-
-    # Make the API call
-    for chunk in ollama.chat(**api_params):
-        yield chunk
-
+from litellm import completion
 
-def process_ollama_tool_stream(
-    stream, tool_map: Dict[str, callable], tools: List[Dict]
-) -> List[Dict]:
-    """Process the Ollama tool use stream"""
-    tool_results = []
-    content = ""
+# import litellm
+# litellm._turn_on_debug()
 
-    # Build tool schema map
-    tool_schemas = {
-        tool["function"]["name"]: tool["function"]["parameters"] for tool in tools
-    }
-
-    def convert_params(tool_name: str, params: dict) -> dict:
-        """Convert parameters to the correct type based on schema"""
-        schema = tool_schemas.get(tool_name, {})
-        properties = schema.get("properties", {})
-
-        converted = {}
-        for key, value in params.items():
-            prop_schema = properties.get(key, {})
-            prop_type = prop_schema.get("type")
-
-            if prop_type == "integer" and isinstance(value, str):
-                try:
-                    converted[key] = int(value)
-                except (ValueError, TypeError):
-                    converted[key] = 0
-            else:
-                converted[key] = value
 
-        return converted
-
-    # Accumulate content
-    for chunk in stream:
-        if chunk.message and chunk.message.content:
-            content += chunk.message.content
-
-    # Process complete JSON objects when done
-    try:
-        # Find all JSON objects in the content
-        json_objects = []
-        current = ""
-        for char in content:
-            current += char
-            if current.count("{") == current.count("}") and current.strip().startswith(
-                "{"
-            ):
-                json_objects.append(current.strip())
-                current = ""
-
-        # Process each JSON object
-        for json_str in json_objects:
-            try:
-                tool_call = json.loads(json_str)
-                tool_name = tool_call.get("name")
-                tool_params = tool_call.get("parameters", {})
-
-                if tool_name in tool_map:
-                    # Convert parameters to correct types
-                    converted_params = convert_params(tool_name, tool_params)
-                    result = tool_map[tool_name](converted_params)
-                    tool_results.append(
-                        {
-                            "tool_name": tool_name,
-                            "tool_input": converted_params,
-                            "tool_result": result,
-                        }
-                    )
-            except Exception as e:
-                tool_results.append({"error": str(e), "partial_json": json_str})
-
-    except Exception as e:
-        tool_results.append({"error": str(e), "content": content})
-
-    return tool_results
-
-
-def get_openai_stream(
+def get_litellm_stream(
     messages: List[Dict[str, str]],
     model: str,
+    provider: str = None,
     npc: Any = None,
     tools: list = None,
     images: List[Dict[str, str]] = None,
     api_key: str = None,
+    api_url: str = None,
     tool_choice: Dict = None,
     **kwargs,
 ) -> Generator:
     """Streams responses from OpenAI, supporting images, tools and yielding raw text chunks."""
-
-    if api_key is None:
-        api_key = os.environ["OPENAI_API_KEY"]
-    client = OpenAI(api_key=api_key)
-
     system_message = get_system_message(npc) if npc else "You are a helpful assistant."
 
     if not messages:
@@ -429,7 +73,26 @@ def get_openai_stream(
             messages.append(last_user_message)
 
     # Prepare API call parameters
-    api_params = {"model": model, "messages": messages, "stream": True, **kwargs}
+    # print("provider", provider)
+    # print("model", model)
+    if provider is not None:
+        model_str = f"{provider}/{model}"
+    else:
+        model_str = model
+
+    api_params = {
+        "model": model_str,
+        "messages": messages,
+        "stream": True,
+    }
+    # print(api_params["model"])
+
+    if api_key is not None and provider == "openai-like":
+        print(api_key)
+        api_params["api_key"] = api_key
+
+    if api_url is not None and provider == "openai-like":
+        api_params["api_url"] = api_url
 
     # Add tools if provided
     if tools:
@@ -438,16 +101,33 @@ def get_openai_stream(
     # Add tool choice if specified
     if tool_choice:
         api_params["tool_choice"] = tool_choice
-
-    stream = client.chat.completions.create(**api_params)
+    if kwargs:
+        for key, value in kwargs.items():
+            if key in [
+                "stream",
+                "stop",
+                "temperature",
+                "top_p",
+                "max_tokens",
+                "max_completion_tokens",
+                "tools",
+                "tool_choice",
+                "extra_headers",
+                "parallel_tool_calls",
+                "response_format",
+                "user",
+            ]:
+                api_params[key] = value
+    # print(api_params)
+    stream = completion(**api_params)
 
     for chunk in stream:
         yield chunk
 
 
-def process_openai_tool_stream(stream, tool_map: Dict[str, callable]) -> List[Dict]:
+def process_litellm_tool_stream(stream, tool_map: Dict[str, callable]) -> List[Dict]:
     """
-    Process the OpenAI tool use stream
+    Process the litellm tool use stream
     """
     final_tool_calls = {}
     tool_results = []
@@ -515,186 +195,39 @@ def process_openai_tool_stream(stream, tool_map: Dict[str, callable]) -> List[Di
     return tool_results
 
 
-def get_openai_like_stream(
-    messages: List[Dict[str, str]],
-    model: str,
-    api_url: str,
-    npc: Any = None,
-    images: List[Dict[str, str]] = None,
-    tools: list = None,
-    api_key: str = None,
-    **kwargs,
-) -> List[Dict[str, str]]:
-    """
-    Function Description:
-        This function generates a conversation using the OpenAI API.
-    Args:
-        messages (List[Dict[str, str]]): The list of messages in the conversation.
-        model (str): The model to use for the conversation.
-    Keyword Args:
-        npc (Any): The NPC object.
-        api_key (str): The API key for accessing the OpenAI API.
-    Returns:
-        List[Dict[str, str]]: The list of messages in the conversation.
-    """
-    if api_key is None:
-        api_key = "dummy"
-    client = OpenAI(api_key=api_key, base_url=api_url)
-
-    system_message = get_system_message(npc) if npc else "You are a helpful assistant."
-
-    if messages is None:
-        messages = []
-
-    # Ensure the system message is at the beginning
-    if not any(msg["role"] == "system" for msg in messages):
-        messages.insert(0, {"role": "system", "content": system_message})
-
-    # messages should already include the user's latest message
-
-    # Make the API call with the messages including the latest user input
-    completion = client.chat.completions.create(
-        model=model, messages=messages, stream=True, **kwargs
-    )
-
-    for chunk in completion:
-        yield chunk
+from typing import List, Dict, Any, Literal
 
 
-def get_deepseek_stream(
-    messages: List[Dict[str, str]],
-    model: str,
-    npc: Any = None,
-    tools: list = None,
-    api_key: str = None,
-    **kwargs,
-) -> List[Dict[str, str]]:
+def generate_tool_schema(
+    name: str,
+    description: str,
+    parameters: Dict[str, Any],
+    required: List[str] = None,
+) -> Dict[str, Any]:
     """
-    Function Description:
-        This function generates a conversation using the Deepseek API.
+    Generate provider-specific function/tool schema from common parameters
+
     Args:
-        messages (List[Dict[str, str]]): The list of messages in the conversation.
-        model (str): The model to use for the conversation.
-    Keyword Args:
-        npc (Any): The NPC object.
-        api_key (str): The API key for accessing the Deepseek API.
-    Returns:
-        List[Dict[str, str]]: The list of messages in the conversation.
+        name: Name of the function
+        description: Description of what the function does
+        parameters: Dict of parameter names and their properties
+        provider: Which provider to generate schema for
+        required: List of required parameter names
     """
-    if api_key is None:
-        api_key = os.environ["DEEPSEEK_API_KEY"]
-    client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
-
-    system_message = get_system_message(npc) if npc else ""
-
-    messages_copy = messages.copy()
-    if messages_copy[0]["role"] != "system":
-        messages_copy.insert(0, {"role": "system", "content": system_message})
-
-    completion = client.chat.completions.create(
-        model=model,
-        messages=messages,
-        tools=tools,
-        stream=True,
-        **kwargs,  # Include any additional keyword arguments
-    )
-
-    for response in completion:
-        yield response
-
-
-def get_gemini_stream(
-    messages: List[Dict[str, str]],
-    model: str,
-    npc: Any = None,
-    tools: list = None,
-    api_key: str = None,
-    **kwargs,
-) -> Generator:
-    """Streams responses from Gemini, supporting tools and yielding chunks."""
-    import google.generativeai as genai
-
-    if api_key is None:
-        api_key = os.environ["GEMINI_API_KEY"]
-
-    # Configure the Gemini API
-    genai.configure(api_key=api_key)
-
-    # Create model instance
-    model = genai.GenerativeModel(model_name=model)
-
-    # Convert all messages to contents list
-    contents = []
-    for msg in messages:
-        if msg["role"] != "system":
-            contents.append(
-                {
-                    "role": "user" if msg["role"] == "user" else "model",
-                    "parts": [{"text": msg["content"]}],
-                }
-            )
-
-    try:
-        # Generate streaming response with full history
-        response = model.generate_content(
-            contents=contents, tools=tools if tools else None, stream=True, **kwargs
-        )
-
-        for chunk in response:
-            yield chunk
-
-    except Exception as e:
-        print(f"Error in Gemini stream: {str(e)}")
-        yield None
-
-
-def process_gemini_tool_stream(
-    stream, tool_map: Dict[str, callable], tools: List[Dict]
-) -> List[Dict]:
-    """Process the Gemini tool stream and execute tools."""
-    tool_results = []
-
-    try:
-        for response in stream:
-            if not response.candidates:
-                continue
-
-            for candidate in response.candidates:
-                if not candidate.content or not candidate.content.parts:
-                    continue
-
-                for part in candidate.content.parts:
-                    if hasattr(part, "function_call"):
-                        try:
-                            tool_name = part.function_call.name
-                            # Convert MapComposite to dict
-                            tool_args = dict(part.function_call.args)
-
-                            if tool_name in tool_map:
-                                result = tool_map[tool_name](tool_args)
-                                tool_results.append(
-                                    {
-                                        "tool_name": tool_name,
-                                        "tool_input": tool_args,
-                                        "tool_result": result,
-                                    }
-                                )
-                        except Exception as e:
-                            tool_results.append(
-                                {
-                                    "error": str(e),
-                                    "tool_name": (
-                                        tool_name
-                                        if "tool_name" in locals()
-                                        else "unknown"
-                                    ),
-                                    "tool_input": (
-                                        tool_args if "tool_args" in locals() else None
-                                    ),
-                                }
-                            )
-
-    except Exception as e:
-        tool_results.append({"error": f"Stream processing error: {str(e)}"})
+    if required is None:
+        required = []
 
-    return tool_results
+    return {
+        "type": "function",
+        "function": {
+            "name": name,
+            "description": description,
+            "parameters": {
+                "type": "object",
+                "properties": parameters,
+                "required": required,
+                "additionalProperties": False,
+            },
+            "strict": True,
+        },
+    }