diff --git a/memoryos-pypi/__init__.py b/memoryos-pypi/__init__.py index b97e620..c7a78ca 100644 --- a/memoryos-pypi/__init__.py +++ b/memoryos-pypi/__init__.py @@ -1,3 +1,4 @@ from .memoryos import Memoryos +from .utils import LiteLLMClient -__all__ = ['Memoryos'] \ No newline at end of file +__all__ = ['Memoryos', 'LiteLLMClient'] \ No newline at end of file diff --git a/memoryos-pypi/memoryos.py b/memoryos-pypi/memoryos.py index 24ebc4c..f4ae5e3 100644 --- a/memoryos-pypi/memoryos.py +++ b/memoryos-pypi/memoryos.py @@ -5,7 +5,7 @@ # 修改为绝对导入 try: # 尝试相对导入（当作为包使用时） - from .utils import OpenAIClient, get_timestamp, generate_id, gpt_user_profile_analysis, gpt_knowledge_extraction, ensure_directory_exists + from .utils import OpenAIClient, LiteLLMClient, get_timestamp, generate_id, gpt_user_profile_analysis, gpt_knowledge_extraction, ensure_directory_exists from . import prompts from .short_term import ShortTermMemory from .mid_term import MidTermMemory, compute_segment_heat # For H_THRESHOLD logic @@ -14,7 +14,7 @@ from .retriever import Retriever except ImportError: # 回退到绝对导入（当作为独立模块使用时） - from utils import OpenAIClient, get_timestamp, generate_id, gpt_user_profile_analysis, gpt_knowledge_extraction, ensure_directory_exists + from utils import OpenAIClient, LiteLLMClient, get_timestamp, generate_id, gpt_user_profile_analysis, gpt_knowledge_extraction, ensure_directory_exists import prompts from short_term import ShortTermMemory from mid_term import MidTermMemory, compute_segment_heat # For H_THRESHOLD logic @@ -27,11 +27,11 @@ DEFAULT_ASSISTANT_ID = "default_assistant_profile" class Memoryos: - def __init__(self, user_id: str, - openai_api_key: str, - data_storage_path: str, - openai_base_url: str = None, - assistant_id: str = DEFAULT_ASSISTANT_ID, + def __init__(self, user_id: str, + openai_api_key: str = None, + data_storage_path: str = ".", + openai_base_url: str = None, + assistant_id: str = DEFAULT_ASSISTANT_ID, short_term_capacity=10, mid_term_capacity=2000, long_term_knowledge_capacity=100, @@ -40,7 +40,10 @@ def __init__(self, user_id: str, mid_term_similarity_threshold=0.6, llm_model="gpt-4o-mini", embedding_model_name: str = "all-MiniLM-L6-v2", - embedding_model_kwargs: dict = None + embedding_model_kwargs: dict = None, + provider: str = "openai", + api_key: str = None, + **provider_kwargs ): self.user_id = user_id self.assistant_id = assistant_id @@ -64,8 +67,20 @@ def __init__(self, user_id: str, print(f"Using unified LLM model: {self.llm_model}") print(f"Using embedding model: {self.embedding_model_name} with kwargs: {self.embedding_model_kwargs}") - # Initialize OpenAI Client - self.client = OpenAIClient(api_key=openai_api_key, base_url=openai_base_url) + # Initialize LLM Client + resolved_key = api_key or openai_api_key + self.provider = provider.lower() + if self.provider == "litellm": + self.client = LiteLLMClient( + api_key=resolved_key, + model=self.llm_model, + **provider_kwargs + ) + else: + self.client = OpenAIClient( + api_key=resolved_key or "", + base_url=openai_base_url + ) # Define file paths for user-specific data self.user_data_dir = os.path.join(self.data_storage_path, "users", self.user_id) diff --git a/memoryos-pypi/requirements.txt b/memoryos-pypi/requirements.txt index 64c0394..2140c9d 100644 --- a/memoryos-pypi/requirements.txt +++ b/memoryos-pypi/requirements.txt @@ -8,6 +8,7 @@ FlagEmbedding>=1.2.9 # For BGE-M3 model support faiss-gpu>=1.7.0,<2.0.0 httpx[socks] openai +litellm>=1.80.0 # Web framework (for demo) flask>=2.0.0,<3.0.0 diff --git a/memoryos-pypi/tests/test_litellm.py b/memoryos-pypi/tests/test_litellm.py new file mode 100644 index 0000000..0075525 --- /dev/null +++ b/memoryos-pypi/tests/test_litellm.py @@ -0,0 +1,268 @@ +"""Comprehensive unit tests for the LiteLLMClient class. + +These tests mock litellm.completion and litellm exception classes so +that litellm does NOT need to be installed in the test environment. +""" + +import os +import sys +import types +import unittest +from unittest.mock import MagicMock, patch + + +# --------------------------------------------------------------------------- +# Bootstrap a fake ``litellm`` package so ``LiteLLMClient`` can be imported +# without the real library installed. +# --------------------------------------------------------------------------- +def _build_fake_litellm(): + """Create minimal stub modules for litellm + litellm.exceptions.""" + + fake_litellm = types.ModuleType("litellm") + fake_exceptions = types.ModuleType("litellm.exceptions") + + # Base exception all litellm errors derive from + class _LiteLLMError(Exception): + def __init__(self, message="", model=None, llm_provider=None, response=None): + self.message = message + self.model = model + self.llm_provider = llm_provider + self.response = response + super().__init__(message) + + class AuthenticationError(_LiteLLMError): + pass + + class RateLimitError(_LiteLLMError): + pass + + class APIConnectionError(_LiteLLMError): + pass + + class Timeout(_LiteLLMError): + pass + + class NotFoundError(_LiteLLMError): + pass + + class BadRequestError(_LiteLLMError): + pass + + class ContextWindowExceededError(_LiteLLMError): + pass + + # Attach to the exceptions sub-module + for cls in ( + AuthenticationError, + RateLimitError, + APIConnectionError, + Timeout, + NotFoundError, + BadRequestError, + ContextWindowExceededError, + ): + setattr(fake_exceptions, cls.__name__, cls) + setattr(fake_litellm, cls.__name__, cls) + + fake_litellm.exceptions = fake_exceptions + fake_litellm.completion = MagicMock() # default; overridden per test + + return fake_litellm, fake_exceptions + + +_fake_litellm, _fake_exceptions = _build_fake_litellm() + +# Inject the fake modules before importing the code under test +sys.modules["litellm"] = _fake_litellm +sys.modules["litellm.exceptions"] = _fake_exceptions + +# Add the memoryos-pypi directory to sys.path so ``utils`` can be +# imported directly. Remove the parent package's __init__.py from +# the import chain to avoid triggering ``from .memoryos import ...`` +# which requires a full package context. +_PACKAGE_DIR = os.path.join(os.path.dirname(__file__), os.pardir) +_PACKAGE_DIR = os.path.abspath(_PACKAGE_DIR) +sys.path.insert(0, _PACKAGE_DIR) + +# Prevent pytest from importing memoryos-pypi/__init__.py which does +# a relative import that fails outside of a proper install. +sys.modules.setdefault("memoryos", types.ModuleType("memoryos")) + +from utils import LiteLLMClient # noqa: E402 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _make_response(content="Hello"): + """Build a mock object shaped like a litellm completion response.""" + message = MagicMock() + message.content = content + choice = MagicMock() + choice.message = message + response = MagicMock() + response.choices = [choice] + return response + + +SAMPLE_MESSAGES = [{"role": "user", "content": "Hi"}] + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +class TestLiteLLMClientHappyPath(unittest.TestCase): + """1. Happy path - valid response.""" + + @patch("litellm.completion") + def test_returns_content(self, mock_completion): + mock_completion.return_value = _make_response("Hello") + client = LiteLLMClient(api_key="sk-test", model="gpt-4o-mini") + result = client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + self.assertEqual(result, "Hello") + + +class TestLiteLLMClientAuthError(unittest.TestCase): + """2. Invalid / expired API key.""" + + @patch("litellm.completion") + def test_authentication_error(self, mock_completion): + mock_completion.side_effect = _fake_exceptions.AuthenticationError( + "Invalid API key" + ) + client = LiteLLMClient(api_key="bad-key", model="gpt-4o-mini") + result = client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + self.assertIn("Authentication", result) + + +class TestLiteLLMClientNotFound(unittest.TestCase): + """3. Model not found.""" + + @patch("litellm.completion") + def test_not_found_error(self, mock_completion): + mock_completion.side_effect = _fake_exceptions.NotFoundError( + "Model not found" + ) + client = LiteLLMClient(api_key="sk-test", model="nonexistent-model") + result = client.chat_completion( + model="nonexistent-model", messages=SAMPLE_MESSAGES + ) + self.assertIn("not found", result.lower()) + + +class TestLiteLLMClientRateLimit(unittest.TestCase): + """4. Rate limit (429).""" + + @patch("litellm.completion") + def test_rate_limit_error(self, mock_completion): + mock_completion.side_effect = _fake_exceptions.RateLimitError( + "Rate limit exceeded" + ) + client = LiteLLMClient(api_key="sk-test", model="gpt-4o-mini") + result = client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + self.assertIn("Rate limit", result) + + +class TestLiteLLMClientContextWindow(unittest.TestCase): + """5. Token limit exceeded.""" + + @patch("litellm.completion") + def test_context_window_exceeded(self, mock_completion): + mock_completion.side_effect = _fake_exceptions.ContextWindowExceededError( + "Context window exceeded" + ) + client = LiteLLMClient(api_key="sk-test", model="gpt-4o-mini") + result = client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + self.assertIn("context window", result.lower()) + + +class TestLiteLLMClientEmptyResponse(unittest.TestCase): + """6. Empty / null response.""" + + @patch("litellm.completion") + def test_none_content_returns_error(self, mock_completion): + mock_completion.return_value = _make_response(content=None) + client = LiteLLMClient(api_key="sk-test", model="gpt-4o-mini") + result = client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + self.assertIn("empty response", result.lower()) + + +class TestLiteLLMClientTimeout(unittest.TestCase): + """7. Timeout.""" + + @patch("litellm.completion") + def test_timeout_error(self, mock_completion): + mock_completion.side_effect = _fake_exceptions.Timeout("Request timed out") + client = LiteLLMClient(api_key="sk-test", model="gpt-4o-mini") + result = client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + self.assertIn("timed out", result.lower()) + + +class TestLiteLLMClientConnectionError(unittest.TestCase): + """8. Connection error.""" + + @patch("litellm.completion") + def test_connection_error(self, mock_completion): + mock_completion.side_effect = _fake_exceptions.APIConnectionError( + "Connection refused" + ) + client = LiteLLMClient(api_key="sk-test", model="gpt-4o-mini") + result = client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + self.assertIn("connect", result.lower()) + + +class TestLiteLLMClientBadRequest(unittest.TestCase): + """9. Bad request.""" + + @patch("litellm.completion") + def test_bad_request_error(self, mock_completion): + mock_completion.side_effect = _fake_exceptions.BadRequestError( + "Invalid request" + ) + client = LiteLLMClient(api_key="sk-test", model="gpt-4o-mini") + result = client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + self.assertIn("Bad request", result) + + +class TestLiteLLMClientDropParams(unittest.TestCase): + """10. drop_params=True is passed.""" + + @patch("litellm.completion") + def test_drop_params_true(self, mock_completion): + mock_completion.return_value = _make_response("ok") + client = LiteLLMClient(api_key="sk-test", model="gpt-4o-mini") + client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + call_kwargs = mock_completion.call_args[1] + self.assertTrue(call_kwargs.get("drop_params"), "drop_params should be True") + + +class TestLiteLLMClientAPIKeyForwarding(unittest.TestCase): + """11. API key forwarding.""" + + @patch("litellm.completion") + def test_api_key_forwarded(self, mock_completion): + mock_completion.return_value = _make_response("ok") + client = LiteLLMClient(api_key="test-key", model="gpt-4o-mini") + client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + call_kwargs = mock_completion.call_args[1] + self.assertEqual(call_kwargs.get("api_key"), "test-key") + + +class TestLiteLLMClientAPIKeyOmitted(unittest.TestCase): + """12. API key omitted when None (env-var fallback).""" + + @patch("litellm.completion") + def test_api_key_not_in_kwargs_when_none(self, mock_completion): + mock_completion.return_value = _make_response("ok") + client = LiteLLMClient(api_key=None, model="gpt-4o-mini") + client.chat_completion(model="gpt-4o-mini", messages=SAMPLE_MESSAGES) + call_kwargs = mock_completion.call_args[1] + self.assertNotIn( + "api_key", + call_kwargs, + "api_key should NOT be in call kwargs when client api_key is None", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/memoryos-pypi/utils.py b/memoryos-pypi/utils.py index 40d95ff..410ccd2 100644 --- a/memoryos-pypi/utils.py +++ b/memoryos-pypi/utils.py @@ -97,6 +97,119 @@ def shutdown(self): """关闭线程池""" self.executor.shutdown(wait=True) + +class LiteLLMClient: + """LLM client powered by LiteLLM, providing access to 100+ providers + (OpenAI, Anthropic, Google, Azure, AWS Bedrock, etc.) through a unified interface. + + Requires: pip install litellm + """ + def __init__(self, api_key=None, model=None, max_workers=5, **kwargs): + try: + import litellm # noqa: F401 + except ImportError: + raise ImportError( + "LiteLLM is required for this provider. " + "Install it with: pip install litellm" + ) + self.api_key = api_key + self.default_model = model + self.extra_kwargs = kwargs + self.executor = ThreadPoolExecutor(max_workers=max_workers) + self._lock = threading.Lock() + + def chat_completion(self, model, messages, temperature=0.7, max_tokens=2000): + import litellm + + effective_model = model or self.default_model + print(f"Calling LiteLLM API. Model: {effective_model}") + + call_kwargs = { + "model": effective_model, + "messages": messages, + "temperature": temperature, + "max_tokens": max_tokens, + "drop_params": True, + } + if self.api_key: + call_kwargs["api_key"] = self.api_key + call_kwargs.update(self.extra_kwargs) + + try: + from litellm.exceptions import ( + AuthenticationError, + RateLimitError, + APIConnectionError, + Timeout, + NotFoundError, + BadRequestError, + ContextWindowExceededError, + ) + except ImportError: + pass + + try: + response = litellm.completion(**call_kwargs) + content = response.choices[0].message.content + if content is None: + return "Error: LLM returned empty response." + raw_content = content.strip() + cleaned_content = clean_reasoning_model_output(raw_content) + return cleaned_content + except AuthenticationError as e: + print(f"LiteLLM auth error: {e}") + return "Error: Authentication failed. Check your API key." + except RateLimitError as e: + print(f"LiteLLM rate limit error: {e}") + return "Error: Rate limit exceeded. Please retry later." + except APIConnectionError as e: + print(f"LiteLLM connection error: {e}") + return "Error: Could not connect to the LLM provider." + except Timeout as e: + print(f"LiteLLM timeout error: {e}") + return "Error: Request timed out." + except NotFoundError as e: + print(f"LiteLLM not found error: {e}") + return "Error: Model or endpoint not found." + except ContextWindowExceededError as e: + print(f"LiteLLM context window exceeded: {e}") + return "Error: Input exceeds the model's context window." + except BadRequestError as e: + print(f"LiteLLM bad request error: {e}") + return "Error: Bad request. Check your input parameters." + except Exception as e: + print(f"Error calling LiteLLM API: {e}") + return "Error: Could not get response from LLM." + + def chat_completion_async(self, model, messages, temperature=0.7, max_tokens=2000): + return self.executor.submit(self.chat_completion, model, messages, temperature, max_tokens) + + def batch_chat_completion(self, requests): + futures = [] + for req in requests: + future = self.chat_completion_async( + model=req.get("model", self.default_model), + messages=req["messages"], + temperature=req.get("temperature", 0.7), + max_tokens=req.get("max_tokens", 2000) + ) + futures.append(future) + + results = [] + for future in as_completed(futures): + try: + result = future.result() + results.append(result) + except Exception as e: + print(f"Error in batch completion: {e}") + results.append("Error: Could not get response from LLM.") + + return results + + def shutdown(self): + self.executor.shutdown(wait=True) + + # ---- Parallel Processing Utilities ---- def run_parallel_tasks(tasks, max_workers=3): """ diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..7b001b9 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +addopts = --import-mode=importlib