powermem 0.1.0__py3-none-any.whl

powermem/integrations/llm/ollama.py

@@ -0,0 +1,106 @@
+from typing import Dict, List, Optional, Union
+
+from powermem.integrations.llm import LLMBase
+from powermem.integrations.llm.config.base import BaseLLMConfig
+from powermem.integrations.llm.config.ollama import OllamaConfig
+
+try:
+    from ollama import Client
+except ImportError:
+    raise ImportError("The 'ollama' library is required. Please install it using 'pip install ollama'.")
+
+
+class OllamaLLM(LLMBase):
+    def __init__(self, config: Optional[Union[BaseLLMConfig, OllamaConfig, Dict]] = None):
+        # Convert to OllamaConfig if needed
+        if config is None:
+            config = OllamaConfig()
+        elif isinstance(config, dict):
+            config = OllamaConfig(**config)
+        elif isinstance(config, BaseLLMConfig) and not isinstance(config, OllamaConfig):
+            # Convert BaseLLMConfig to OllamaConfig
+            config = OllamaConfig(
+                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 = "llama3.1:70b"
+
+        self.client = Client(host=self.config.ollama_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["message"]["content"] if isinstance(response, dict) else response.message.content,
+                "tool_calls": [],
+            }
+
+            # Ollama doesn't support tool calls in the same way, so we return the content
+            return processed_response
+        else:
+            # Handle both dict and object responses
+            if isinstance(response, dict):
+                return response["message"]["content"]
+            else:
+                return response.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 Ollama.
+
+        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 Ollama-specific parameters.
+
+        Returns:
+            str: The generated response.
+        """
+        # Build parameters for Ollama
+        params = {
+            "model": self.config.model,
+            "messages": messages,
+        }
+
+        # Add options for Ollama (temperature, num_predict, top_p)
+        options = {
+            "temperature": self.config.temperature,
+            "num_predict": self.config.max_tokens,
+            "top_p": self.config.top_p,
+        }
+        params["options"] = options
+
+        # Remove OpenAI-specific parameters that Ollama doesn't support
+        params.pop("max_tokens", None)  # Ollama uses different parameter names
+
+        response = self.client.chat(**params)
+        return self._parse_response(response, tools)

powermem/integrations/llm/openai.py

@@ -0,0 +1,166 @@
+import json
+import logging
+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.openai import OpenAIConfig
+from powermem.utils.utils import extract_json
+
+
+class OpenAILLM(LLMBase):
+    def __init__(self, config: Optional[Union[BaseLLMConfig, OpenAIConfig, Dict]] = None):
+        # Convert to OpenAIConfig if needed
+        if config is None:
+            config = OpenAIConfig()
+        elif isinstance(config, dict):
+            config = OpenAIConfig(**config)
+        elif isinstance(config, BaseLLMConfig) and not isinstance(config, OpenAIConfig):
+            # Convert BaseLLMConfig to OpenAIConfig
+            config = OpenAIConfig(
+                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 = "gpt-4o-mini"
+
+        if os.environ.get("OPENROUTER_API_KEY"):  # Use OpenRouter
+            self.client = OpenAI(
+                api_key=os.environ.get("OPENROUTER_API_KEY"),
+                base_url=self.config.openrouter_base_url
+                or os.getenv("OPENROUTER_API_BASE")
+                or "https://openrouter.ai/api/v1",
+            )
+        else:
+            api_key = self.config.api_key or os.getenv("OPENAI_API_KEY")
+            base_url = self.config.openai_base_url or os.getenv("OPENAI_BASE_URL") or "https://api.openai.com/v1"
+
+            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:
+                    # Extract and validate arguments
+                    arguments_str = extract_json(tool_call.function.arguments)
+
+                    # Check if arguments are empty or whitespace only
+                    if not arguments_str or arguments_str.strip() == "":
+                        logging.warning(
+                            f"Tool call '{tool_call.function.name}' has empty arguments. Skipping this tool call."
+                        )
+                        continue
+
+                    # Try to parse JSON with error handling
+                    try:
+                        arguments = json.loads(arguments_str)
+                    except json.JSONDecodeError as e:
+                        logging.error(
+                            f"Failed to parse tool call arguments for '{tool_call.function.name}': "
+                            f"{arguments_str[:100]}... Error: {e}"
+                        )
+                        continue
+
+                    processed_response["tool_calls"].append(
+                        {
+                            "name": tool_call.function.name,
+                            "arguments": 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 JSON response based on the given messages using OpenAI.
+
+        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 OpenAI-specific parameters.
+
+        Returns:
+            json: The generated response.
+        """
+        params = self._get_supported_params(messages=messages, **kwargs)
+        
+        params.update({
+            "model": self.config.model,
+            "messages": messages,
+        })
+
+        if os.getenv("OPENROUTER_API_KEY"):
+            openrouter_params = {}
+            if self.config.models:
+                openrouter_params["models"] = self.config.models
+                openrouter_params["route"] = self.config.route
+                params.pop("model")
+
+            if self.config.site_url and self.config.app_name:
+                extra_headers = {
+                    "HTTP-Referer": self.config.site_url,
+                    "X-Title": self.config.app_name,
+                }
+                openrouter_params["extra_headers"] = extra_headers
+
+            params.update(**openrouter_params)
+        
+        else:
+            openai_specific_generation_params = ["store"]
+            for param in openai_specific_generation_params:
+                if hasattr(self.config, param):
+                    params[param] = getattr(self.config, param)
+            
+        if response_format:
+            params["response_format"] = response_format
+        if tools:  # TODO: Remove tools if no issues found with new memory addition logic
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+        response = self.client.chat.completions.create(**params)
+        parsed_response = self._parse_response(response, tools)
+        if self.config.response_callback:
+            try:
+                self.config.response_callback(self, response, params)
+            except Exception as e:
+                # Log error but don't propagate
+                logging.error(f"Error due to callback: {e}")
+                pass
+        return parsed_response

powermem/integrations/llm/openai_structured.py

@@ -0,0 +1,80 @@
+import json
+import os
+from typing import Dict, List, Optional, Any
+
+from openai import OpenAI
+
+from powermem.integrations.llm import LLMBase
+from powermem.integrations.llm.config.base import BaseLLMConfig
+
+
+class OpenAIStructuredLLM(LLMBase):
+    def __init__(self, config: Optional[BaseLLMConfig] = None):
+        super().__init__(config)
+
+        if not self.config.model:
+            self.config.model = "gpt-5"
+
+        api_key = self.config.api_key or os.getenv("OPENAI_API_KEY")
+        base_url = self.config.openai_base_url or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1"
+        self.client = OpenAI(api_key=api_key, base_url=base_url)
+
+    def generate_response(
+            self,
+            messages: List[Dict[str, str]],
+            response_format: Optional[str] = None,
+            tools: Optional[List[Dict]] = None,
+            tool_choice: str = "auto",
+    ) -> str | None | dict[str, str | None | list[Any]]:
+        """
+        Generate a response based on the given messages using OpenAI.
+
+        Args:
+            messages (List[Dict[str, str]]): A list of dictionaries, each containing a 'role' and 'content' key.
+            response_format (Optional[str]): The desired format of the response. Defaults to None.
+
+
+        Returns:
+            str: The generated response.
+        """
+        params = {
+            "model": self.config.model,
+            "messages": messages,
+            "temperature": self.config.temperature,
+        }
+
+        if response_format:
+            params["response_format"] = response_format
+        if tools:
+            params["tools"] = tools
+            params["tool_choice"] = tool_choice
+
+        response = self.client.beta.chat.completions.parse(**params)
+
+        message = response.choices[0].message
+
+        # If tools were used and tool_calls exist, return structured response
+        if tools and message.tool_calls:
+            tool_calls_list = []
+            for tool_call in message.tool_calls:
+                arguments = tool_call.function.arguments
+                # Parse arguments if it's a string
+                if isinstance(arguments, str):
+                    try:
+                        arguments = json.loads(arguments)
+                    except json.JSONDecodeError:
+                        arguments = {}
+
+                tool_calls_list.append({
+                    "id": tool_call.id,
+                    "name": tool_call.function.name,
+                    "arguments": arguments
+                })
+
+            return {
+                "content": message.content,
+                "tool_calls": tool_calls_list
+            }
+
+        # Otherwise return just the content
+        return message.content or ""

powermem/integrations/llm/qwen.py

@@ -0,0 +1,207 @@
+import json
+import logging
+import os
+from typing import Dict, List, Optional, Union
+
+try:
+    from dashscope import Generation
+    from dashscope.api_entities.dashscope_response import DashScopeAPIResponse
+except ImportError:
+    Generation = None
+    DashScopeAPIResponse = None
+
+from powermem.integrations.llm import LLMBase
+from powermem.integrations.llm.config.base import BaseLLMConfig
+from powermem.integrations.llm.config.qwen import QwenConfig
+from powermem.utils.utils import extract_json
+import dashscope
+
+
+class QwenLLM(LLMBase):
+    def __init__(self, config: Optional[Union[BaseLLMConfig, QwenConfig, Dict]] = None):
+        # Check if dashscope is available first
+        try:
+            from dashscope import Generation
+            from dashscope.api_entities.dashscope_response import DashScopeAPIResponse
+        except ImportError:
+            raise ImportError(
+                "DashScope SDK is not installed. Please install it with: pip install dashscope"
+            )
+
+        # Convert to QwenConfig if needed
+        if config is None:
+            config = QwenConfig()
+        elif isinstance(config, dict):
+            config = QwenConfig(**config)
+        elif isinstance(config, BaseLLMConfig) and not isinstance(config, QwenConfig):
+            # Convert BaseLLMConfig to QwenConfig
+            config = QwenConfig(
+                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 = "qwen-turbo"
+
+        # Set API key
+        api_key = self.config.api_key or os.getenv("DASHSCOPE_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "API key is required. Set DASHSCOPE_API_KEY environment variable or pass api_key in config.")
+
+        # Set API key for DashScope SDK
+        dashscope.api_key = api_key
+
+        # Set base URL
+        base_url = self.config.dashscope_base_url or os.getenv(
+            "DASHSCOPE_BASE_URL") or "https://dashscope.aliyuncs.com/api/v1"
+
+        if base_url:
+            os.environ["DASHSCOPE_BASE_URL"] = base_url
+
+    def _get_attr(self, obj, key, default=None):
+        """Unified handling of attribute access for both dicts and objects"""
+        if isinstance(obj, dict):
+            return obj.get(key, default)
+        return getattr(obj, key, default)
+    
+    def _extract_message(self, output):
+        """Extract message object from response output"""
+        # Get default text content
+        text = self._get_attr(output, 'text', '')
+        
+        # Extract message from choices
+        choices = self._get_attr(output, 'choices', [])
+        if choices:
+            choice = choices[0]
+            message = self._get_attr(choice, 'message')
+            return message, text
+        
+        return None, text
+    
+    def _extract_content(self, output):
+        """Extract response content"""
+        message, default_text = self._extract_message(output)
+        if message:
+            return self._get_attr(message, 'content', default_text)
+        return default_text
+    
+    def _extract_tool_calls(self, output):
+        """Extract tool calls from response"""
+        message, _ = self._extract_message(output)
+        if not message:
+            return []
+        
+        tool_calls = self._get_attr(message, 'tool_calls')
+        if not tool_calls:
+            return []
+        
+        processed_calls = []
+        for tool_call in tool_calls:
+            function = self._get_attr(tool_call, 'function', {})
+            name = self._get_attr(function, 'name')
+            arguments = self._get_attr(function, 'arguments', '{}')
+            
+            processed_calls.append({
+                "name": name,
+                "arguments": json.loads(extract_json(arguments)),
+            })
+        
+        return processed_calls
+
+    def _parse_response(self, response: DashScopeAPIResponse, tools: Optional[List[Dict]] = None):
+        """
+        Process the response based on whether tools are used or not.
+
+        Args:
+            response: The raw response from DashScope API.
+            tools: The list of tools provided in the request.
+
+        Returns:
+            str or dict: The processed response.
+        """
+        if response.status_code != 200:
+            raise Exception(f"API request failed with status {response.status_code}: {response.message}")
+
+        content = self._extract_content(response.output)
+        
+        if tools:
+            return {
+                "content": content,
+                "tool_calls": self._extract_tool_calls(response.output),
+            }
+        
+        return 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 Qwen.
+
+        Args:
+            messages (list): List of message dicts containing 'role' and 'content'.
+            response_format (str or object, optional): Format of the response. Defaults to None.
+            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 Qwen-specific parameters.
+
+        Returns:
+            str or dict: The generated response.
+        """
+        params = self._get_supported_params(**kwargs)
+
+        # Prepare generation parameters
+        generation_params = {
+            "model": self.config.model,
+            "messages": messages,
+            "temperature": params.get("temperature", self.config.temperature),
+            "max_tokens": params.get("max_tokens", self.config.max_tokens),
+            "top_p": params.get("top_p", self.config.top_p),
+        }
+
+        # Add Qwen-specific parameters
+        if self.config.enable_search:
+            generation_params["enable_search"] = True
+            if self.config.search_params:
+                generation_params.update(self.config.search_params)
+
+        # Add tools if provided
+        if tools:
+            generation_params["tools"] = tools
+            generation_params["tool_choice"] = tool_choice
+
+        # Add response format if provided
+        if response_format:
+            generation_params["response_format"] = response_format
+
+        try:
+            response = Generation.call(**generation_params)
+            parsed_response = self._parse_response(response, tools)
+
+            if self.config.response_callback:
+                try:
+                    self.config.response_callback(self, response, generation_params)
+                except Exception as e:
+                    # Log error but don't propagate
+                    logging.error(f"Error due to callback: {e}")
+                    pass
+
+            return parsed_response
+        except Exception as e:
+            logging.error(f"Qwen API call failed: {e}")
+            raise