powermem 0.1.0__py3-none-any.whl

powermem/integrations/llm/config/qwen_asr.py

@@ -0,0 +1,46 @@
+from typing import Optional
+
+from powermem.integrations.llm.config.base import BaseLLMConfig
+
+
+class QwenASRConfig(BaseLLMConfig):
+    """
+    Configuration class for Qwen ASR-specific parameters.
+    Inherits from BaseLLMConfig and adds ASR-specific settings.
+    """
+
+    def __init__(
+            self,
+            # Base parameters (only model and api_key are used for ASR)
+            model: Optional[str] = None,
+            api_key: Optional[str] = None,
+            # ASR-specific parameters
+            dashscope_base_url: Optional[str] = None,
+            asr_options: Optional[dict] = None,
+            result_format: str = "message",
+    ):
+        """
+        Initialize Qwen ASR configuration.
+
+        Args:
+            model: Qwen ASR model to use, defaults to "qwen3-asr-flash"
+            api_key: DashScope API key, defaults to None
+            dashscope_base_url: DashScope API base URL, defaults to None
+            asr_options: ASR-specific options (e.g., language, enable_itn), defaults to {"enable_itn": True}
+            result_format: Result format for ASR response, defaults to "message"
+        """
+        # Initialize base parameters with defaults (ASR doesn't use these parameters)
+        super().__init__(
+            model=model,
+            api_key=api_key,
+        )
+
+        # ASR-specific parameters
+        self.dashscope_base_url = dashscope_base_url
+        # Default asr_options with enable_itn enabled
+        if asr_options is None:
+            self.asr_options = {"enable_itn": True}
+        else:
+            self.asr_options = asr_options
+        self.result_format = result_format
+

powermem/integrations/llm/config/vllm.py

@@ -0,0 +1,56 @@
+from typing import Optional
+
+from powermem.integrations.llm.config.base import BaseLLMConfig
+
+
+class VllmConfig(BaseLLMConfig):
+    """
+    Configuration class for vLLM-specific parameters.
+    Inherits from BaseLLMConfig and adds vLLM-specific settings.
+    """
+
+    def __init__(
+            self,
+            # Base parameters
+            model: Optional[str] = None,
+            temperature: float = 0.1,
+            api_key: Optional[str] = None,
+            max_tokens: int = 2000,
+            top_p: float = 0.1,
+            top_k: int = 1,
+            enable_vision: bool = False,
+            vision_details: Optional[str] = "auto",
+            http_client_proxies: Optional[dict] = None,
+            # vLLM-specific parameters
+            vllm_base_url: Optional[str] = None,
+    ):
+        """
+        Initialize vLLM configuration.
+
+        Args:
+            model: vLLM model to use, defaults to None
+            temperature: Controls randomness, defaults to 0.1
+            api_key: vLLM API key, defaults to None
+            max_tokens: Maximum tokens to generate, defaults to 2000
+            top_p: Nucleus sampling parameter, defaults to 0.1
+            top_k: Top-k sampling parameter, defaults to 1
+            enable_vision: Enable vision capabilities, defaults to False
+            vision_details: Vision detail level, defaults to "auto"
+            http_client_proxies: HTTP client proxy settings, defaults to None
+            vllm_base_url: vLLM base URL, defaults to None
+        """
+        # Initialize base parameters
+        super().__init__(
+            model=model,
+            temperature=temperature,
+            api_key=api_key,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            top_k=top_k,
+            enable_vision=enable_vision,
+            vision_details=vision_details,
+            http_client_proxies=http_client_proxies,
+        )
+
+        # vLLM-specific parameters
+        self.vllm_base_url = vllm_base_url or "http://localhost:8000/v1"

powermem/integrations/llm/configs.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class LLMConfig(BaseModel):
+    provider: str = Field(description="Provider of the LLM (e.g., 'ollama', 'openai')", default="openai")
+    config: Optional[dict] = Field(description="Configuration for the specific LLM", default={})
+
+    @field_validator("config")
+    def validate_config(cls, v, info):
+        provider = info.data.get("provider")
+        if provider in (
+            "openai",
+            "ollama",
+            "anthropic",
+            "openai_structured",
+            "gemini",
+            "deepseek",
+            "vllm",
+            "langchain",
+            "qwen",
+        ):
+            return v
+        else:
+            raise ValueError(f"Unsupported LLM provider: {provider}")

powermem/integrations/llm/deepseek.py

@@ -0,0 +1,106 @@
+import json
+import os
+from typing import Dict, List, Optional, Union
+
+from openai import OpenAI
+from powermem.integrations.llm import LLMBase
+from powermem.integrations.llm.config.base import BaseLLMConfig
+from powermem.integrations.llm.config.deepseek import DeepSeekConfig
+from powermem.utils.utils import extract_json
+
+
+class DeepSeekLLM(LLMBase):
+    def __init__(self, config: Optional[Union[BaseLLMConfig, DeepSeekConfig, Dict]] = None):
+        # Convert to DeepSeekConfig if needed
+        if config is None:
+            config = DeepSeekConfig()
+        elif isinstance(config, dict):
+            config = DeepSeekConfig(**config)
+        elif isinstance(config, BaseLLMConfig) and not isinstance(config, DeepSeekConfig):
+            # Convert BaseLLMConfig to DeepSeekConfig
+            config = DeepSeekConfig(
+                model=config.model,
+                temperature=config.temperature,
+                api_key=config.api_key,
+                max_tokens=config.max_tokens,
+                top_p=config.top_p,
+                top_k=config.top_k,
+                enable_vision=config.enable_vision,
+                vision_details=config.vision_details,
+                http_client_proxies=config.http_client,
+            )
+
+        super().__init__(config)
+
+        if not self.config.model:
+            self.config.model = "deepseek-chat"
+
+        api_key = self.config.api_key or os.getenv("DEEPSEEK_API_KEY")
+        base_url = self.config.deepseek_base_url or os.getenv("DEEPSEEK_API_BASE") or "https://api.deepseek.com"
+        self.client = OpenAI(api_key=api_key, base_url=base_url)
+
+    def _parse_response(self, response, tools):
+        """
+        Process the response based on whether tools are used or not.
+
+        Args:
+            response: The raw response from API.
+            tools: The list of tools provided in the request.
+
+        Returns:
+            str or dict: The processed response.
+        """
+        if tools:
+            processed_response = {
+                "content": response.choices[0].message.content,
+                "tool_calls": [],
+            }
+
+            if response.choices[0].message.tool_calls:
+                for tool_call in response.choices[0].message.tool_calls:
+                    processed_response["tool_calls"].append(
+                        {
+                            "name": tool_call.function.name,
+                            "arguments": json.loads(extract_json(tool_call.function.arguments)),
+                        }
+                    )
+
+            return processed_response
+        else:
+            return response.choices[0].message.content
+
+    def generate_response(
+            self,
+            messages: List[Dict[str, str]],
+            response_format=None,
+            tools: Optional[List[Dict]] = None,
+            tool_choice: str = "auto",
+            **kwargs,
+    ):
+        """
+        Generate a response based on the given messages using DeepSeek.
+
+        Args:
+            messages (list): List of message dicts containing 'role' and 'content'.
+            response_format (str or object, optional): Format of the response. Defaults to "text".
+            tools (list, optional): List of tools that the model can call. Defaults to None.
+            tool_choice (str, optional): Tool choice method. Defaults to "auto".
+            **kwargs: Additional DeepSeek-specific parameters.
+
+        Returns:
+            str: The generated response.
+        """
+        params = self._get_supported_params(messages=messages, **kwargs)
+        params.update(
+            {
+                "model": self.config.model,
+                "messages": messages,
+            }
+        )
+
+        if tools:
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+
+        response = self.client.chat.completions.create(**params)
+        return self._parse_response(response, tools)

powermem/integrations/llm/factory.py

@@ -0,0 +1,118 @@
+import importlib
+from typing import Dict, Optional, Union
+
+from powermem.integrations.llm.config.anthropic import AnthropicConfig
+from powermem.integrations.llm.config.base import BaseLLMConfig
+from powermem.integrations.llm.config.deepseek import DeepSeekConfig
+from powermem.integrations.llm.config.ollama import OllamaConfig
+from powermem.integrations.llm.config.openai import OpenAIConfig
+from powermem.integrations.llm.config.qwen import QwenConfig
+from powermem.integrations.llm.config.qwen_asr import QwenASRConfig
+from powermem.integrations.llm.config.vllm import VllmConfig
+
+
+def load_class(class_type):
+    module_path, class_name = class_type.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)
+
+
+class LLMFactory:
+    """
+    Factory for creating LLM instances with appropriate configurations.
+    Supports both old-style BaseLLMConfig and new provider-specific configs.
+    """
+
+    # Provider mappings with their config classes
+    provider_to_class = {
+        "ollama": ("powermem.integrations.llm.ollama.OllamaLLM", OllamaConfig),
+        "openai": ("powermem.integrations.llm.openai.OpenAILLM", OpenAIConfig),
+        "openai_structured": ("powermem.integrations.llm.openai_structured.OpenAIStructuredLLM", OpenAIConfig),
+        "anthropic": ("powermem.integrations.llm.anthropic.AnthropicLLM", AnthropicConfig),
+        "gemini": ("powermem.integrations.llm.gemini.GeminiLLM", BaseLLMConfig),
+        "deepseek": ("powermem.integrations.llm.deepseek.DeepSeekLLM", DeepSeekConfig),
+        "vllm": ("powermem.integrations.llm.vllm.VllmLLM", VllmConfig),
+        "langchain": ("powermem.integrations.llm.langchain.LangchainLLM", BaseLLMConfig),
+        "qwen": ("powermem.integrations.llm.qwen.QwenLLM", QwenConfig),
+        "qwen_asr": ("powermem.integrations.llm.qwen_asr.QwenASR", QwenASRConfig),
+    }
+
+    @classmethod
+    def create(cls, provider_name: str, config: Optional[Union[BaseLLMConfig, Dict]] = None, **kwargs):
+        """
+        Create an LLM instance with the appropriate configuration.
+
+        Args:
+            provider_name (str): The provider name (e.g., 'openai', 'anthropic')
+            config: Configuration object or dict. If None, will create default config
+            **kwargs: Additional configuration parameters
+
+        Returns:
+            Configured LLM instance
+
+        Raises:
+            ValueError: If provider is not supported
+        """
+        if provider_name not in cls.provider_to_class:
+            raise ValueError(f"Unsupported Llm provider: {provider_name}")
+
+        class_type, config_class = cls.provider_to_class[provider_name]
+        llm_class = load_class(class_type)
+
+        # Handle configuration
+        if config is None:
+            # Create default config with kwargs
+            config = config_class(**kwargs)
+        elif isinstance(config, dict):
+            # Merge dict config with kwargs
+            config.update(kwargs)
+            config = config_class(**config)
+        elif isinstance(config, BaseLLMConfig):
+            # Convert base config to provider-specific config if needed
+            if config_class != BaseLLMConfig:
+                # Convert to provider-specific config
+                config_dict = {
+                    "model": config.model,
+                    "temperature": config.temperature,
+                    "api_key": config.api_key,
+                    "max_tokens": config.max_tokens,
+                    "top_p": config.top_p,
+                    "top_k": config.top_k,
+                    "enable_vision": config.enable_vision,
+                    "vision_details": config.vision_details,
+                    "http_client_proxies": config.http_client,
+                }
+                config_dict.update(kwargs)
+                config = config_class(**config_dict)
+            else:
+                # Use base config as-is
+                pass
+        else:
+            # Assume it's already the correct config type
+            pass
+
+        return llm_class(config)
+
+    @classmethod
+    def register_provider(cls, name: str, class_path: str, config_class=None):
+        """
+        Register a new provider.
+
+        Args:
+            name (str): Provider name
+            class_path (str): Full path to LLM class
+            config_class: Configuration class for the provider (defaults to BaseLLMConfig)
+        """
+        if config_class is None:
+            config_class = BaseLLMConfig
+        cls.provider_to_class[name] = (class_path, config_class)
+
+    @classmethod
+    def get_supported_providers(cls) -> list:
+        """
+        Get list of supported providers.
+
+        Returns:
+            list: List of supported provider names
+        """
+        return list(cls.provider_to_class.keys())

powermem/integrations/llm/gemini.py

@@ -0,0 +1,201 @@
+import os
+from typing import Dict, List, Optional
+
+from powermem.integrations.llm import LLMBase
+from powermem.integrations.llm.config.base import BaseLLMConfig
+
+try:
+    from google import genai
+    from google.genai import types
+except ImportError:
+    raise ImportError("The 'google-genai' library is required. Please install it using 'pip install google-genai'.")
+
+
+class GeminiLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLLMConfig] = None):
+        super().__init__(config)
+
+        if not self.config.model:
+            self.config.model = "gemini-2.0-flash"
+
+        api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
+        self.client = genai.Client(api_key=api_key)
+
+    def _parse_response(self, response, tools):
+        """
+        Process the response based on whether tools are used or not.
+
+        Args:
+            response: The raw response from API.
+            tools: The list of tools provided in the request.
+
+        Returns:
+            str or dict: The processed response.
+        """
+        if tools:
+            processed_response = {
+                "content": None,
+                "tool_calls": [],
+            }
+
+            # Extract content from the first candidate
+            if response.candidates and response.candidates[0].content.parts:
+                for part in response.candidates[0].content.parts:
+                    if hasattr(part, "text") and part.text:
+                        processed_response["content"] = part.text
+                        break
+
+            # Extract function calls
+            if response.candidates and response.candidates[0].content.parts:
+                for part in response.candidates[0].content.parts:
+                    if hasattr(part, "function_call") and part.function_call:
+                        fn = part.function_call
+                        processed_response["tool_calls"].append(
+                            {
+                                "name": fn.name,
+                                "arguments": dict(fn.args) if fn.args else {},
+                            }
+                        )
+
+            return processed_response
+        else:
+            if response.candidates and response.candidates[0].content.parts:
+                for part in response.candidates[0].content.parts:
+                    if hasattr(part, "text") and part.text:
+                        return part.text
+            return ""
+
+    def _reformat_messages(self, messages: List[Dict[str, str]]):
+        """
+        Reformat messages for Gemini.
+
+        Args:
+            messages: The list of messages provided in the request.
+
+        Returns:
+            tuple: (system_instruction, contents_list)
+        """
+        system_instruction = None
+        contents = []
+
+        for message in messages:
+            if message["role"] == "system":
+                system_instruction = message["content"]
+            else:
+                content = types.Content(
+                    parts=[types.Part(text=message["content"])],
+                    role=message["role"],
+                )
+                contents.append(content)
+
+        return system_instruction, contents
+
+    def _reformat_tools(self, tools: Optional[List[Dict]]):
+        """
+        Reformat tools for Gemini.
+
+        Args:
+            tools: The list of tools provided in the request.
+
+        Returns:
+            list: The list of tools in the required format.
+        """
+
+        def remove_additional_properties(data):
+            """Recursively removes 'additionalProperties' from nested dictionaries."""
+            if isinstance(data, dict):
+                filtered_dict = {
+                    key: remove_additional_properties(value)
+                    for key, value in data.items()
+                    if not (key == "additionalProperties")
+                }
+                return filtered_dict
+            else:
+                return data
+
+        if tools:
+            function_declarations = []
+            for tool in tools:
+                func = tool["function"].copy()
+                cleaned_func = remove_additional_properties(func)
+
+                function_declaration = types.FunctionDeclaration(
+                    name=cleaned_func["name"],
+                    description=cleaned_func.get("description", ""),
+                    parameters=cleaned_func.get("parameters", {}),
+                )
+                function_declarations.append(function_declaration)
+
+            tool_obj = types.Tool(function_declarations=function_declarations)
+            return [tool_obj]
+        else:
+            return None
+
+    def generate_response(
+            self,
+            messages: List[Dict[str, str]],
+            response_format=None,
+            tools: Optional[List[Dict]] = None,
+            tool_choice: str = "auto",
+    ):
+        """
+        Generate a response based on the given messages using Gemini.
+
+        Args:
+            messages (list): List of message dicts containing 'role' and 'content'.
+            response_format (str or object, optional): Format for the response. Defaults to "text".
+            tools (list, optional): List of tools that the model can call. Defaults to None.
+            tool_choice (str, optional): Tool choice method. Defaults to "auto".
+
+        Returns:
+            str: The generated response.
+        """
+
+        # Extract system instruction and reformat messages
+        system_instruction, contents = self._reformat_messages(messages)
+
+        # Prepare generation config
+        config_params = {
+            "temperature": self.config.temperature,
+            "max_output_tokens": self.config.max_tokens,
+            "top_p": self.config.top_p,
+        }
+
+        # Add system instruction to config if present
+        if system_instruction:
+            config_params["system_instruction"] = system_instruction
+
+        if response_format is not None and response_format["type"] == "json_object":
+            config_params["response_mime_type"] = "application/json"
+            if "schema" in response_format:
+                config_params["response_schema"] = response_format["schema"]
+
+        if tools:
+            formatted_tools = self._reformat_tools(tools)
+            config_params["tools"] = formatted_tools
+
+            if tool_choice:
+                if tool_choice == "auto":
+                    mode = types.FunctionCallingConfigMode.AUTO
+                elif tool_choice == "any":
+                    mode = types.FunctionCallingConfigMode.ANY
+                else:
+                    mode = types.FunctionCallingConfigMode.NONE
+
+                tool_config = types.ToolConfig(
+                    function_calling_config=types.FunctionCallingConfig(
+                        mode=mode,
+                        allowed_function_names=(
+                            [tool["function"]["name"] for tool in tools] if tool_choice == "any" else None
+                        ),
+                    )
+                )
+                config_params["tool_config"] = tool_config
+
+        generation_config = types.GenerateContentConfig(**config_params)
+
+        response = self.client.models.generate_content(
+            model=self.config.model, contents=contents, config=generation_config
+        )
+
+        return self._parse_response(response, tools)

powermem/integrations/llm/langchain.py

@@ -0,0 +1,65 @@
+from typing import Dict, List, Optional
+
+from powermem.integrations.llm import LLMBase
+from powermem.integrations.llm.config.base import BaseLLMConfig
+
+try:
+    from langchain.chat_models.base import BaseChatModel
+except ImportError:
+    raise ImportError("langchain is not installed. Please install it using `pip install langchain`")
+
+
+class LangchainLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLLMConfig] = None):
+        super().__init__(config)
+
+        if self.config.model is None:
+            raise ValueError("`model` parameter is required")
+
+        if not isinstance(self.config.model, BaseChatModel):
+            raise ValueError("`model` must be an instance of BaseChatModel")
+
+        self.langchain_model = self.config.model
+
+    def generate_response(
+            self,
+            messages: List[Dict[str, str]],
+            response_format=None,
+            tools: Optional[List[Dict]] = None,
+            tool_choice: str = "auto",
+    ):
+        """
+        Generate a response based on the given messages using langchain_community.
+
+        Args:
+            messages (list): List of message dicts containing 'role' and 'content'.
+            response_format (str or object, optional): Format of the response. Not used in Langchain.
+            tools (list, optional): List of tools that the model can call. Not used in Langchain.
+            tool_choice (str, optional): Tool choice method. Not used in Langchain.
+
+        Returns:
+            str: The generated response.
+        """
+        try:
+            # Convert the messages to LangChain's tuple format
+            langchain_messages = []
+            for message in messages:
+                role = message["role"]
+                content = message["content"]
+
+                if role == "system":
+                    langchain_messages.append(("system", content))
+                elif role == "user":
+                    langchain_messages.append(("human", content))
+                elif role == "assistant":
+                    langchain_messages.append(("ai", content))
+
+            if not langchain_messages:
+                raise ValueError("No valid messages found in the messages list")
+
+            ai_message = self.langchain_model.invoke(langchain_messages)
+
+            return ai_message.content
+
+        except Exception as e:
+            raise Exception(f"Error generating response using langchain model: {str(e)}")