memmachine-server 0.3.3__tar.gz → 0.3.4__tar.gz

{memmachine_server-0.3.3/src/memmachine_server.egg-info → memmachine_server-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: memmachine-server
-Version: 0.3.3
+Version: 0.3.4
 Summary: MemMachine Server - The complete MemMachine memory system server with episodic and profile memory
 Author-email: MemMachine <noreply@memmachine.ai>
 License-Expression: Apache-2.0

{memmachine_server-0.3.3 → memmachine_server-0.3.4}/server_tests/memmachine_server/common/configuration/test_reranker_conf.py

@@ -7,6 +7,7 @@ from pydantic import SecretStr, ValidationError
 from memmachine_server.common.configuration.reranker_conf import (
     AmazonBedrockRerankerConf,
     BM25RerankerConf,
+    CohereRerankerConf,
     CrossEncoderRerankerConf,
     EmbedderRerankerConf,
     IdentityRerankerConf,
@@ -47,6 +48,18 @@ def amazon_bedrock_reranker_conf() -> dict:
     }
 
 
+@pytest.fixture
+def cohere_reranker_conf() -> dict:
+    return {
+        "provider": "cohere",
+        "config": {
+            "cohere_key": "test-cohere-key",
+            "model": "rerank-english-v3.0",
+            "base_url": "http://localhost:8000",
+        },
+    }
+
+
 @pytest.fixture
 def cross_encoder_reranker_conf() -> dict:
     return {
@@ -83,6 +96,7 @@ def full_reranker_input(
     bm25_reranker_conf,
     identity_reranker_conf,
     amazon_bedrock_reranker_conf,
+    cohere_reranker_conf,
     cross_encoder_reranker_conf,
     embedder_reranker_conf,
     rrf_hybrid_reranker_conf,
@@ -93,6 +107,7 @@ def full_reranker_input(
             "id_ranker_id": identity_reranker_conf,
             "bm_ranker_id": bm25_reranker_conf,
             "aws_reranker_id": amazon_bedrock_reranker_conf,
+            "cohere_reranker_id": cohere_reranker_conf,
             "cross_encoder_id": cross_encoder_reranker_conf,
             "embedder_id": embedder_reranker_conf,
         },
@@ -121,6 +136,13 @@ def test_valid_amazon_bedrock_reranker_conf(amazon_bedrock_reranker_conf):
     assert conf.model_id == "amazon.rerank-v1:0"
 
 
+def test_valid_cohere_reranker_conf(cohere_reranker_conf):
+    conf = CohereRerankerConf(**cohere_reranker_conf["config"])
+    assert conf.cohere_key == SecretStr("test-cohere-key")
+    assert conf.model == "rerank-english-v3.0"
+    assert conf.base_url == "http://localhost:8000"
+
+
 def test_valid_cross_encoder_reranker_conf(cross_encoder_reranker_conf):
     conf = CrossEncoderRerankerConf(**cross_encoder_reranker_conf["config"])
     assert conf.model_name == "cross-encoder/qnli-electra-base"
@@ -150,6 +172,8 @@ def test_full_reranker_conf(full_reranker_input):
 
     assert "aws_reranker_id" in conf.amazon_bedrock
     assert conf.amazon_bedrock["aws_reranker_id"].region == "us-west-2"
+    assert "cohere_reranker_id" in conf.cohere
+    assert conf.cohere["cohere_reranker_id"].base_url == "http://localhost:8000"
 
     assert "cross_encoder_id" in conf.cross_encoder
     assert (
@@ -182,3 +206,12 @@ def test_missing_required_field_in_bedrock_reranker():
         AmazonBedrockRerankerConf(**config)
     assert "missing" in str(exc_info.value)
     assert "aws_secret_access_key" in str(exc_info.value)
+
+
+def test_invalid_cohere_base_url():
+    with pytest.raises(ValidationError, match="Invalid base URL"):
+        CohereRerankerConf(
+            cohere_key=SecretStr("test-cohere-key"),
+            model="rerank-english-v3.0",
+            base_url="localhost:8000",
+        )

{memmachine_server-0.3.3 → memmachine_server-0.3.4}/server_tests/memmachine_server/common/language_model/test_openai_chat_completions_language_model.py

@@ -1,5 +1,7 @@
 """Unit tests for OpenAIChatCompletionsLanguageModel."""
 
+from dataclasses import dataclass
+from typing import Any, cast
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
@@ -9,7 +11,7 @@ from openai.types import chat as openai_chat
 from openai.types.chat.chat_completion_message_function_tool_call import (
     Function as ToolCallFunction,
 )
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 
 from memmachine_server.common.data_types import ExternalServiceAPIError
 from memmachine_server.common.language_model.openai_chat_completions_language_model import (
@@ -19,6 +21,38 @@ from memmachine_server.common.language_model.openai_chat_completions_language_mo
 from memmachine_server.common.metrics_factory import MetricsFactory
 
 
+class FakeAsyncStream:
+    def __init__(self, chunks):
+        self._chunks = iter(chunks)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return next(self._chunks)
+        except StopIteration as exc:
+            raise StopAsyncIteration from exc
+
+
+@dataclass
+class FakeReasoningChunk:
+    type: str = "response.reasoning.delta"
+
+
+class ParsedResponse(BaseModel):
+    answer: str
+    confidence: float
+
+
+def make_chat_completion(data: dict[str, Any]) -> openai_chat.ChatCompletion:
+    return openai_chat.ChatCompletion.model_validate(data)
+
+
+def make_chat_completion_chunk(data: dict[str, Any]) -> openai_chat.ChatCompletionChunk:
+    return openai_chat.ChatCompletionChunk.model_validate(data)
+
+
 @pytest.fixture
 def mock_metrics_factory():
     """Fixture for a mocked MetricsFactory."""
@@ -213,6 +247,339 @@ async def test_generate_response_success(mock_async_openai, minimal_config):
         {"role": "system", "content": "System prompt"},
         {"role": "user", "content": "User prompt"},
     ]
+    assert call_args.kwargs["store"] is False
+
+
+@pytest.mark.asyncio
+async def test_generate_parsed_response_success(mock_async_openai, minimal_config):
+    """Test a non-streamed structured response is parsed correctly."""
+    mock_response = make_chat_completion(
+        {
+            "id": "completion-1",
+            "object": "chat.completion",
+            "created": 1,
+            "model": "test-model",
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": "stop",
+                    "message": {
+                        "role": "assistant",
+                        "content": '{"answer":"ok","confidence":0.9}',
+                        "refusal": None,
+                    },
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 7,
+                "completion_tokens": 5,
+                "total_tokens": 12,
+            },
+        }
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = mock_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    parsed = await lm.generate_parsed_response(
+        ParsedResponse,
+        system_prompt="System prompt",
+        user_prompt="User prompt",
+    )
+
+    assert parsed == ParsedResponse(answer="ok", confidence=0.9)
+    call_args = mock_client.chat.completions.create.call_args
+    assert call_args.kwargs["model"] == "test-model"
+    assert call_args.kwargs["messages"] == [
+        {"role": "system", "content": "System prompt"},
+        {"role": "user", "content": "User prompt"},
+    ]
+    assert call_args.kwargs["response_format"]["type"] == "json_schema"
+
+
+@pytest.mark.asyncio
+async def test_generate_parsed_response_from_stream(
+    mock_async_openai,
+    minimal_config,
+):
+    """Test a streamed structured response is aggregated and validated."""
+    streamed_response = FakeAsyncStream(
+        [
+            FakeReasoningChunk(),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-1",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "role": "assistant",
+                                "content": '{"answer":"ok",',
+                            },
+                            "finish_reason": None,
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-2",
+                    "object": "chat.completion.chunk",
+                    "created": 2,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": '"confidence":0.9}'},
+                            "finish_reason": "stop",
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+        ]
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = streamed_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    parsed = await lm.generate_parsed_response(ParsedResponse)
+
+    assert parsed == ParsedResponse(answer="ok", confidence=0.9)
+
+
+@pytest.mark.asyncio
+async def test_generate_parsed_response_from_stream_discards_reasoning_delta_chunks(
+    mock_async_openai,
+    minimal_config,
+):
+    """Test streamed structured responses ignore chunks with reasoning_content."""
+    reasoning_chunk = make_chat_completion_chunk(
+        {
+            "id": "chunk-0",
+            "object": "chat.completion.chunk",
+            "created": 0,
+            "model": "test-model",
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"content": None},
+                    "finish_reason": None,
+                    "logprobs": None,
+                }
+            ],
+            "usage": None,
+        }
+    )
+    cast(Any, reasoning_chunk.choices[0].delta).reasoning_content = "chain-of-thought"
+
+    streamed_response = FakeAsyncStream(
+        [
+            reasoning_chunk,
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-1",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "role": "assistant",
+                                "content": '{"answer":"ok",',
+                            },
+                            "finish_reason": None,
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-2",
+                    "object": "chat.completion.chunk",
+                    "created": 2,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": '"confidence":0.9}'},
+                            "finish_reason": "stop",
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+        ]
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = streamed_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    parsed = await lm.generate_parsed_response(ParsedResponse)
+
+    assert parsed == ParsedResponse(answer="ok", confidence=0.9)
+
+
+@pytest.mark.asyncio
+async def test_generate_response_streamed_chat_completion_chunks(
+    mock_async_openai,
+    minimal_config,
+):
+    """Test a streamed chat completion is aggregated into a single response."""
+    streamed_response = FakeAsyncStream(
+        [
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-1",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"role": "assistant", "content": "Hello"},
+                            "finish_reason": None,
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-2",
+                    "object": "chat.completion.chunk",
+                    "created": 2,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": ", world!"},
+                            "finish_reason": "stop",
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": {
+                        "prompt_tokens": 10,
+                        "completion_tokens": 4,
+                        "total_tokens": 14,
+                    },
+                }
+            ),
+        ]
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = streamed_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    (
+        content,
+        tool_calls,
+        input_tokens,
+        output_tokens,
+    ) = await lm.generate_response_with_token_usage()
+
+    assert content == "Hello, world!"
+    assert tool_calls == []
+    assert input_tokens == 10
+    assert output_tokens == 4
+
+
+@pytest.mark.asyncio
+async def test_generate_response_streamed_tool_calls_and_discards_reasoning_chunks(
+    mock_async_openai,
+    minimal_config,
+):
+    """Test streamed tool calls are reconstructed while reasoning chunks are ignored."""
+    streamed_response = FakeAsyncStream(
+        [
+            FakeReasoningChunk(),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-1",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "tool_calls": [
+                                    {
+                                        "index": 0,
+                                        "id": "call_123",
+                                        "type": "function",
+                                        "function": {
+                                            "name": "get_weather",
+                                            "arguments": '{"location"',
+                                        },
+                                    }
+                                ]
+                            },
+                            "finish_reason": None,
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-2",
+                    "object": "chat.completion.chunk",
+                    "created": 2,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "tool_calls": [
+                                    {
+                                        "index": 0,
+                                        "function": {"arguments": ': "Boston"}'},
+                                    }
+                                ]
+                            },
+                            "finish_reason": "tool_calls",
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+        ]
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = streamed_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    content, tool_calls = await lm.generate_response()
+
+    assert content == ""
+    assert tool_calls == [
+        {
+            "call_id": "call_123",
+            "function": {
+                "name": "get_weather",
+                "arguments": {"location": "Boston"},
+            },
+        }
+    ]
 
 
 @pytest.mark.asyncio

{memmachine_server-0.3.3 → memmachine_server-0.3.4}/server_tests/memmachine_server/common/language_model/test_openai_responses_language_model.py

@@ -184,6 +184,7 @@ async def test_generate_response_success(mock_async_openai, minimal_config):
         {"role": "system", "content": "System prompt"},
         {"role": "user", "content": "User prompt"},
     ]
+    assert call_args.kwargs["store"] is False
 
 
 @pytest.mark.asyncio
@@ -219,6 +220,8 @@ async def test_generate_response_with_tool_calls(mock_async_openai, minimal_conf
             },
         },
     ]
+    call_args = mock_client.responses.create.call_args
+    assert call_args.kwargs["store"] is False
 
 
 @pytest.mark.asyncio

{memmachine_server-0.3.3 → memmachine_server-0.3.4}/server_tests/memmachine_server/common/resource_manager/test_reranker_manager.py

@@ -34,8 +34,9 @@ def mock_conf():
         bm25={"bm_ranker_id": BM25RerankerConf(tokenizer="simple")},
         cohere={
             "cohere_reranker_id": CohereRerankerConf(
-                cohere_key=SecretStr("<COHERE_API_KEY>"),
+                cohere_key=SecretStr("test-cohere-key"),
                 model="rerank-english-v3.0",
+                base_url="http://localhost:8000",
             ),
         },
         cross_encoder={
@@ -132,6 +133,72 @@ async def test_build_cohere_rerankers(reranker_manager):
     assert reranker is not None
 
 
+@pytest.mark.asyncio
+async def test_build_cohere_reranker_passes_base_url(monkeypatch):
+    captured_kwargs = {}
+
+    class FakeCohereClient:
+        pass
+
+    def fake_client_v2(**kwargs):
+        captured_kwargs.update(kwargs)
+        return FakeCohereClient()
+
+    monkeypatch.setattr("cohere.ClientV2", fake_client_v2)
+
+    conf = RerankersConf(
+        cohere={
+            "cohere_reranker_id": CohereRerankerConf(
+                cohere_key=SecretStr("test-cohere-key"),
+                model="rerank-english-v3.0",
+                base_url="http://localhost:8000",
+            ),
+        },
+    )
+    reranker_manager = RerankerManager(
+        conf=conf,
+        embedder_factory=cast(EmbedderFactory, FakeEmbedderFactory()),
+    )
+
+    await reranker_manager.get_reranker("cohere_reranker_id")
+
+    assert captured_kwargs == {
+        "api_key": "test-cohere-key",
+        "base_url": "http://localhost:8000",
+    }
+
+
+@pytest.mark.asyncio
+async def test_build_cohere_reranker_omits_base_url_when_not_configured(monkeypatch):
+    captured_kwargs = {}
+
+    class FakeCohereClient:
+        pass
+
+    def fake_client_v2(**kwargs):
+        captured_kwargs.update(kwargs)
+        return FakeCohereClient()
+
+    monkeypatch.setattr("cohere.ClientV2", fake_client_v2)
+
+    conf = RerankersConf(
+        cohere={
+            "cohere_reranker_id": CohereRerankerConf(
+                cohere_key=SecretStr("test-cohere-key"),
+                model="rerank-english-v3.0",
+            ),
+        },
+    )
+    reranker_manager = RerankerManager(
+        conf=conf,
+        embedder_factory=cast(EmbedderFactory, FakeEmbedderFactory()),
+    )
+
+    await reranker_manager.get_reranker("cohere_reranker_id")
+
+    assert captured_kwargs == {"api_key": "test-cohere-key"}
+
+
 @requires_sentence_transformers
 @pytest.mark.asyncio
 async def test_build_cross_encoder_rerankers(reranker_manager):