MemoryOS 2.0.3__py3-none-any.whl

memos/llms/hf_singleton.py

@@ -0,0 +1,114 @@
+import threading
+
+from typing import ClassVar
+
+from memos.configs.llm import HFLLMConfig
+from memos.llms.hf import HFLLM
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class HFSingletonLLM(HFLLM):
+    """
+    Singleton version of HFLLM that prevents multiple loading of the same model.
+    This class inherits from HFLLM and adds singleton behavior.
+    """
+
+    _instances: ClassVar[dict[str, "HFSingletonLLM"]] = {}
+    _lock: ClassVar[threading.Lock] = threading.Lock()
+
+    def __new__(cls, config: HFLLMConfig):
+        """
+        Singleton pattern implementation.
+        Returns existing instance if config already exists, otherwise creates new one.
+        """
+        config_key = cls._get_config_key(config)
+
+        if config_key in cls._instances:
+            logger.debug(f"Reusing existing HF model: {config.model_name_or_path}")
+            return cls._instances[config_key]
+
+        with cls._lock:
+            # Double-check pattern to prevent race conditions
+            if config_key in cls._instances:
+                logger.debug(f"Reusing existing HF model: {config.model_name_or_path}")
+                return cls._instances[config_key]
+
+            logger.info(f"Creating new HF model: {config.model_name_or_path}")
+            instance = super().__new__(cls)
+            cls._instances[config_key] = instance
+            return instance
+
+    def __init__(self, config: HFLLMConfig):
+        """
+        Initialize the singleton HFLLM instance.
+        Only initializes if this is a new instance.
+        """
+        # Check if already initialized
+        if hasattr(self, "_initialized"):
+            return
+
+        # Call parent constructor
+        super().__init__(config)
+        self._initialized = True
+
+    @classmethod
+    def _get_config_key(cls, config: HFLLMConfig) -> str:
+        """
+        Generate a unique key for the HF model configuration.
+
+        Args:
+            config: The HFLLM configuration
+
+        Returns:
+            A unique string key representing the configuration
+        """
+        # Create a unique key based on model path and key parameters
+        key_parts = [config.model_name_or_path]
+        return "|".join(key_parts)
+
+    @classmethod
+    def get_instance_count(cls) -> int:
+        """
+        Get the number of unique HF model instances currently managed.
+
+        Returns:
+            Number of HF model instances
+        """
+        return len(cls._instances)
+
+    @classmethod
+    def get_instance_info(cls) -> dict[str, str]:
+        """
+        Get information about all managed HF model instances.
+
+        Returns:
+            Dictionary mapping config keys to model paths
+        """
+        return {key: instance.config.model_name_or_path for key, instance in cls._instances.items()}
+
+    @classmethod
+    def clear_all(cls) -> None:
+        """
+        Clear all HF model instances from memory.
+        This should be used carefully as it will force reloading of models.
+        """
+        with cls._lock:
+            cls._instances.clear()
+            logger.info("All HF model instances cleared from singleton manager")
+
+
+# Convenience function to get singleton manager info
+def get_hf_singleton_info() -> dict[str, int]:
+    """
+    Get information about the HF singleton manager.
+
+    Returns:
+        Dictionary with instance count and info
+    """
+    return {
+        "instance_count": HFSingletonLLM.get_instance_count(),
+        "instance_info": HFSingletonLLM.get_instance_info(),
+    }

memos/llms/ollama.py

@@ -0,0 +1,135 @@
+from collections.abc import Generator
+from typing import Any
+
+from ollama import Client, Message
+
+from memos.configs.llm import OllamaLLMConfig
+from memos.llms.base import BaseLLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageList
+
+
+logger = get_logger(__name__)
+
+
+class OllamaLLM(BaseLLM):
+    """Ollama LLM class."""
+
+    def __init__(self, config: OllamaLLMConfig):
+        self.config = config
+        self.api_base = config.api_base
+
+        # Default model if not specified
+        if not self.config.model_name_or_path:
+            self.config.model_name_or_path = "llama3.1:latest"
+
+        # Initialize ollama client
+        self.client = Client(host=self.api_base)
+
+        # Ensure the model exists locally
+        self._ensure_model_exists()
+
+    def _list_models(self) -> list[str]:
+        """
+        List all models available in the Ollama client.
+
+        Returns:
+            List of model names.
+        """
+        local_models = self.client.list()["models"]
+        return [model.model for model in local_models]
+
+    def _ensure_model_exists(self):
+        """
+        Ensure the specified model exists locally. If not, pull it from Ollama.
+        """
+        try:
+            local_models = self._list_models()
+            if self.config.model_name_or_path not in local_models:
+                logger.warning(
+                    f"Model {self.config.model_name_or_path} not found locally. Pulling from Ollama..."
+                )
+                self.client.pull(self.config.model_name_or_path)
+        except Exception as e:
+            logger.warning(f"Could not verify model existence: {e}")
+
+    def generate(self, messages: MessageList, **kwargs) -> Any:
+        """
+        Generate a response from Ollama LLM.
+
+        Args:
+            messages: List of message dicts containing 'role' and 'content'.
+
+        Returns:
+            str: The generated response.
+        """
+        response = self.client.chat(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            options={
+                "temperature": kwargs.get("temperature", self.config.temperature),
+                "num_predict": kwargs.get("max_tokens", self.config.max_tokens),
+                "top_p": kwargs.get("top_p", self.config.top_p),
+                "top_k": kwargs.get("top_k", self.config.top_k),
+            },
+            think=self.config.enable_thinking,
+            tools=kwargs.get("tools"),
+        )
+        logger.info(f"Raw response from Ollama: {response.model_dump_json()}")
+        tool_calls = getattr(response.message, "tool_calls", None)
+        if isinstance(tool_calls, list) and len(tool_calls) > 0:
+            return self.tool_call_parser(tool_calls)
+
+        str_thinking = (
+            f"<think>{response.message.thinking}</think>"
+            if hasattr(response.message, "thinking")
+            else ""
+        )
+        str_response = response.message.content
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(str_response)
+        else:
+            return str_thinking + str_response
+
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        if kwargs.get("tools"):
+            logger.info("stream api not support tools")
+            return
+
+        response = self.client.chat(
+            model=kwargs.get("model_name_or_path", self.config.model_name_or_path),
+            messages=messages,
+            options={
+                "temperature": kwargs.get("temperature", self.config.temperature),
+                "num_predict": kwargs.get("max_tokens", self.config.max_tokens),
+                "top_p": kwargs.get("top_p", self.config.top_p),
+                "top_k": kwargs.get("top_k", self.config.top_k),
+            },
+            think=self.config.enable_thinking,
+            stream=True,
+        )
+        # Streaming chunks of text
+        reasoning_started = False
+        for chunk in response:
+            if hasattr(chunk.message, "thinking") and chunk.message.thinking:
+                if not reasoning_started and not self.config.remove_think_prefix:
+                    yield "<think>"
+                    reasoning_started = True
+                yield chunk.message.thinking
+
+            if hasattr(chunk.message, "content") and chunk.message.content:
+                if reasoning_started and not self.config.remove_think_prefix:
+                    yield "</think>"
+                    reasoning_started = False
+                yield chunk.message.content
+
+    def tool_call_parser(self, tool_calls: list[Message.ToolCall]) -> list[dict]:
+        """Parse tool calls from OpenAI response."""
+        return [
+            {
+                "function_name": tool_call.function.name,
+                "arguments": tool_call.function.arguments,
+            }
+            for tool_call in tool_calls
+        ]

memos/llms/openai.py

@@ -0,0 +1,222 @@
+import json
+import time
+
+from collections.abc import Generator
+
+import openai
+
+from openai._types import NOT_GIVEN
+from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall
+
+from memos.configs.llm import AzureLLMConfig, OpenAILLMConfig
+from memos.llms.base import BaseLLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageList
+from memos.utils import timed_with_status
+
+
+logger = get_logger(__name__)
+
+
+class OpenAILLM(BaseLLM):
+    """OpenAI LLM class via openai.chat.completions.create."""
+
+    def __init__(self, config: OpenAILLMConfig):
+        self.config = config
+        self.client = openai.Client(
+            api_key=config.api_key, base_url=config.api_base, default_headers=config.default_headers
+        )
+        logger.info("OpenAI LLM instance initialized")
+
+    @timed_with_status(
+        log_prefix="OpenAI LLM",
+        log_extra_args=lambda self, messages, **kwargs: {
+            "model_name_or_path": kwargs.get("model_name_or_path", self.config.model_name_or_path),
+            "messages": messages,
+        },
+    )
+    def generate(self, messages: MessageList, **kwargs) -> str:
+        """Generate a response from OpenAI LLM, optionally overriding generation params."""
+        request_body = {
+            "model": kwargs.get("model_name_or_path", self.config.model_name_or_path),
+            "messages": messages,
+            "temperature": kwargs.get("temperature", self.config.temperature),
+            "max_tokens": kwargs.get("max_tokens", self.config.max_tokens),
+            "top_p": kwargs.get("top_p", self.config.top_p),
+            "extra_body": kwargs.get("extra_body", self.config.extra_body),
+            "tools": kwargs.get("tools", NOT_GIVEN),
+        }
+        start_time = time.perf_counter()
+        logger.info(f"OpenAI LLM Request body: {request_body}")
+
+        response = self.client.chat.completions.create(**request_body)
+
+        cost_time = time.perf_counter() - start_time
+        logger.info(
+            f"Request body: {request_body}, Response from OpenAI: {response.model_dump_json()}, Cost time: {cost_time}"
+        )
+
+        if not response.choices:
+            logger.warning("OpenAI response has no choices")
+            return ""
+
+        tool_calls = getattr(response.choices[0].message, "tool_calls", None)
+        if isinstance(tool_calls, list) and len(tool_calls) > 0:
+            return self.tool_call_parser(tool_calls)
+        response_content = response.choices[0].message.content
+        reasoning_content = getattr(response.choices[0].message, "reasoning_content", None)
+        if isinstance(reasoning_content, str) and reasoning_content:
+            reasoning_content = f"<think>{reasoning_content}</think>"
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(response_content)
+        if reasoning_content:
+            return reasoning_content + (response_content or "")
+        return response_content or ""
+
+    @timed_with_status(
+        log_prefix="OpenAI LLM Stream",
+        log_extra_args=lambda self, messages, **kwargs: {
+            "model_name_or_path": self.config.model_name_or_path
+        },
+    )
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        """Stream response from OpenAI LLM with optional reasoning support."""
+        if kwargs.get("tools"):
+            logger.info("stream api not support tools")
+            return
+
+        request_body = {
+            "model": self.config.model_name_or_path,
+            "messages": messages,
+            "stream": True,
+            "temperature": kwargs.get("temperature", self.config.temperature),
+            "max_tokens": kwargs.get("max_tokens", self.config.max_tokens),
+            "top_p": kwargs.get("top_p", self.config.top_p),
+            "extra_body": kwargs.get("extra_body", self.config.extra_body),
+            "tools": kwargs.get("tools", NOT_GIVEN),
+        }
+
+        logger.info(f"OpenAI LLM Stream Request body: {request_body}")
+        response = self.client.chat.completions.create(**request_body)
+
+        reasoning_started = False
+
+        for chunk in response:
+            if not chunk.choices:
+                continue
+            delta = chunk.choices[0].delta
+
+            # Support for custom 'reasoning_content' (if present in OpenAI-compatible models like Qwen, DeepSeek)
+            if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+                if not reasoning_started and not self.config.remove_think_prefix:
+                    yield "<think>"
+                    reasoning_started = True
+                yield delta.reasoning_content
+            elif hasattr(delta, "content") and delta.content:
+                if reasoning_started and not self.config.remove_think_prefix:
+                    yield "</think>"
+                    reasoning_started = False
+                yield delta.content
+
+        # Ensure we close the <think> block if not already done
+        if reasoning_started and not self.config.remove_think_prefix:
+            yield "</think>"
+
+    def tool_call_parser(self, tool_calls: list[ChatCompletionMessageToolCall]) -> list[dict]:
+        """Parse tool calls from OpenAI response."""
+        return [
+            {
+                "tool_call_id": tool_call.id,
+                "function_name": tool_call.function.name,
+                "arguments": json.loads(tool_call.function.arguments),
+            }
+            for tool_call in tool_calls
+        ]
+
+
+class AzureLLM(BaseLLM):
+    """Azure OpenAI LLM class with singleton pattern."""
+
+    def __init__(self, config: AzureLLMConfig):
+        self.config = config
+        self.client = openai.AzureOpenAI(
+            azure_endpoint=config.base_url,
+            api_version=config.api_version,
+            api_key=config.api_key,
+        )
+        logger.info("Azure LLM instance initialized")
+
+    def generate(self, messages: MessageList, **kwargs) -> str:
+        """Generate a response from Azure OpenAI LLM."""
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            temperature=kwargs.get("temperature", self.config.temperature),
+            max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
+            top_p=kwargs.get("top_p", self.config.top_p),
+            tools=kwargs.get("tools", NOT_GIVEN),
+            extra_body=kwargs.get("extra_body", self.config.extra_body),
+        )
+        logger.info(f"Response from Azure OpenAI: {response.model_dump_json()}")
+        if not response.choices:
+            logger.warning("Azure OpenAI response has no choices")
+            return ""
+
+        if response.choices[0].message.tool_calls:
+            return self.tool_call_parser(response.choices[0].message.tool_calls)
+        response_content = response.choices[0].message.content
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(response_content)
+        else:
+            return response_content or ""
+
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        """Stream response from Azure OpenAI LLM with optional reasoning support."""
+        if kwargs.get("tools"):
+            logger.info("stream api not support tools")
+            return
+
+        response = self.client.chat.completions.create(
+            model=self.config.model_name_or_path,
+            messages=messages,
+            stream=True,
+            temperature=kwargs.get("temperature", self.config.temperature),
+            max_tokens=kwargs.get("max_tokens", self.config.max_tokens),
+            top_p=kwargs.get("top_p", self.config.top_p),
+            extra_body=kwargs.get("extra_body", self.config.extra_body),
+        )
+
+        reasoning_started = False
+
+        for chunk in response:
+            if not chunk.choices:
+                continue
+            delta = chunk.choices[0].delta
+
+            # Support for custom 'reasoning_content' (if present in OpenAI-compatible models like Qwen, DeepSeek)
+            if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+                if not reasoning_started and not self.config.remove_think_prefix:
+                    yield "<think>"
+                    reasoning_started = True
+                yield delta.reasoning_content
+            elif hasattr(delta, "content") and delta.content:
+                if reasoning_started and not self.config.remove_think_prefix:
+                    yield "</think>"
+                    reasoning_started = False
+                yield delta.content
+
+        # Ensure we close the <think> block if not already done
+        if reasoning_started and not self.config.remove_think_prefix:
+            yield "</think>"
+
+    def tool_call_parser(self, tool_calls: list[ChatCompletionMessageToolCall]) -> list[dict]:
+        """Parse tool calls from OpenAI response."""
+        return [
+            {
+                "tool_call_id": tool_call.id,
+                "function_name": tool_call.function.name,
+                "arguments": json.loads(tool_call.function.arguments),
+            }
+            for tool_call in tool_calls
+        ]

memos/llms/openai_new.py

@@ -0,0 +1,198 @@
+import json
+
+from collections.abc import Generator
+
+import openai
+
+from openai._types import NOT_GIVEN
+from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
+from openai.types.responses.response_reasoning_item import ResponseReasoningItem
+
+from memos.configs.llm import AzureLLMConfig, OpenAILLMConfig
+from memos.llms.base import BaseLLM
+from memos.llms.utils import remove_thinking_tags
+from memos.log import get_logger
+from memos.types import MessageList
+from memos.utils import timed
+
+
+logger = get_logger(__name__)
+
+
+class OpenAIResponsesLLM(BaseLLM):
+    def __init__(self, config: OpenAILLMConfig):
+        self.config = config
+        self.client = openai.Client(
+            api_key=config.api_key, base_url=config.api_base, default_headers=config.default_headers
+        )
+
+    @timed(log=True, log_prefix="OpenAI Responses LLM")
+    def generate(self, messages: MessageList, **kwargs) -> str:
+        response = self.client.responses.create(
+            model=kwargs.get("model_name_or_path", self.config.model_name_or_path),
+            input=messages,
+            temperature=kwargs.get("temperature", self.config.temperature),
+            top_p=kwargs.get("top_p", self.config.top_p),
+            max_output_tokens=kwargs.get("max_tokens", self.config.max_tokens),
+            reasoning={"effort": "low", "summary": "auto"}
+            if self.config.enable_thinking
+            else NOT_GIVEN,
+            tools=kwargs.get("tools", NOT_GIVEN),
+            extra_body=kwargs.get("extra_body", self.config.extra_body),
+        )
+        tool_call_outputs = [
+            item for item in response.output if isinstance(item, ResponseFunctionToolCall)
+        ]
+        if tool_call_outputs:
+            return self.tool_call_parser(tool_call_outputs)
+
+        output_text = getattr(response, "output_text", "")
+        output_reasoning = [
+            item for item in response.output if isinstance(item, ResponseReasoningItem)
+        ]
+        summary = output_reasoning[0].summary
+
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(output_text)
+        if summary:
+            return f"<think>{summary[0].text}</think>" + output_text
+        return output_text
+
+    @timed(log=True, log_prefix="OpenAI Responses LLM")
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        if kwargs.get("tools"):
+            logger.info("stream api not support tools")
+            return
+
+        stream = self.client.responses.create(
+            model=kwargs.get("model_name_or_path", self.config.model_name_or_path),
+            input=messages,
+            temperature=kwargs.get("temperature", self.config.temperature),
+            top_p=kwargs.get("top_p", self.config.top_p),
+            max_output_tokens=kwargs.get("max_tokens", self.config.max_tokens),
+            reasoning={"effort": "low", "summary": "auto"}
+            if self.config.enable_thinking
+            else NOT_GIVEN,
+            extra_body=kwargs.get("extra_body", self.config.extra_body),
+            stream=True,
+        )
+
+        reasoning_started = False
+
+        for event in stream:
+            event_type = getattr(event, "type", "")
+            if event_type in (
+                "response.reasoning.delta",
+                "response.reasoning_summary_text.delta",
+            ) and hasattr(event, "delta"):
+                if not self.config.remove_think_prefix:
+                    if not reasoning_started:
+                        yield "<think>"
+                        reasoning_started = True
+                    yield event.delta
+            elif event_type == "response.output_text.delta" and hasattr(event, "delta"):
+                if reasoning_started and not self.config.remove_think_prefix:
+                    yield "</think>"
+                    reasoning_started = False
+                yield event.delta
+
+        if reasoning_started and not self.config.remove_think_prefix:
+            yield "</think>"
+
+    def tool_call_parser(self, tool_calls: list[ResponseFunctionToolCall]) -> list[dict]:
+        """Parse tool calls from OpenAI response."""
+        return [
+            {
+                "tool_call_id": tool_call.call_id,
+                "function_name": tool_call.name,
+                "arguments": json.loads(tool_call.arguments),
+            }
+            for tool_call in tool_calls
+        ]
+
+
+class AzureResponsesLLM(BaseLLM):
+    def __init__(self, config: AzureLLMConfig):
+        self.config = config
+        self.client = openai.AzureOpenAI(
+            azure_endpoint=config.base_url,
+            api_version=config.api_version,
+            api_key=config.api_key,
+        )
+
+    def generate(self, messages: MessageList, **kwargs) -> str:
+        response = self.client.responses.create(
+            model=self.config.model_name_or_path,
+            input=messages,
+            temperature=kwargs.get("temperature", self.config.temperature),
+            top_p=kwargs.get("top_p", self.config.top_p),
+            max_output_tokens=kwargs.get("max_tokens", self.config.max_tokens),
+            tools=kwargs.get("tools", NOT_GIVEN),
+            extra_body=kwargs.get("extra_body", self.config.extra_body),
+            reasoning={"effort": "low", "summary": "auto"}
+            if self.config.enable_thinking
+            else NOT_GIVEN,
+        )
+
+        output_text = getattr(response, "output_text", "")
+        output_reasoning = [
+            item for item in response.output if isinstance(item, ResponseReasoningItem)
+        ]
+        summary = output_reasoning[0].summary
+
+        if self.config.remove_think_prefix:
+            return remove_thinking_tags(output_text)
+        if summary:
+            return f"<think>{summary[0].text}</think>" + output_text
+        return output_text
+
+    def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
+        if kwargs.get("tools"):
+            logger.info("stream api not support tools")
+            return
+
+        stream = self.client.responses.create(
+            model=self.config.model_name_or_path,
+            input=messages,
+            temperature=kwargs.get("temperature", self.config.temperature),
+            top_p=kwargs.get("top_p", self.config.top_p),
+            max_output_tokens=kwargs.get("max_tokens", self.config.max_tokens),
+            extra_body=kwargs.get("extra_body", self.config.extra_body),
+            stream=True,
+            reasoning={"effort": "low", "summary": "auto"}
+            if self.config.enable_thinking
+            else NOT_GIVEN,
+        )
+
+        reasoning_started = False
+
+        for event in stream:
+            event_type = getattr(event, "type", "")
+            if event_type in (
+                "response.reasoning.delta",
+                "response.reasoning_summary_text.delta",
+            ) and hasattr(event, "delta"):
+                if not self.config.remove_think_prefix:
+                    if not reasoning_started:
+                        yield "<think>"
+                        reasoning_started = True
+                    yield event.delta
+            elif event_type == "response.output_text.delta" and hasattr(event, "delta"):
+                if reasoning_started and not self.config.remove_think_prefix:
+                    yield "</think>"
+                    reasoning_started = False
+                yield event.delta
+
+        if reasoning_started and not self.config.remove_think_prefix:
+            yield "</think>"
+
+    def tool_call_parser(self, tool_calls: list[ResponseFunctionToolCall]) -> list[dict]:
+        """Parse tool calls from OpenAI response."""
+        return [
+            {
+                "tool_call_id": tool_call.call_id,
+                "function_name": tool_call.name,
+                "arguments": json.loads(tool_call.arguments),
+            }
+            for tool_call in tool_calls
+        ]

memos/llms/qwen.py

@@ -0,0 +1,13 @@
+from memos.configs.llm import QwenLLMConfig
+from memos.llms.openai import OpenAILLM
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class QwenLLM(OpenAILLM):
+    """Qwen (DashScope) LLM class via OpenAI-compatible API."""
+
+    def __init__(self, config: QwenLLMConfig):
+        super().__init__(config)