120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import httpx
|
|
|
|
from javis.agent.middleware import MemoryMiddleware, TurnContext
|
|
from javis.core.config import Config, PathsConfig
|
|
from javis.llm.hermes_client import HermesClient
|
|
from javis.llm.ollama_client import ChatResponse
|
|
|
|
|
|
class HermesClientTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_chat_uses_bearer_auth_and_openai_payload(self) -> None:
|
|
config = Config()
|
|
config.memory.provider = "hermes"
|
|
config.memory.base_url = "http://hermes.test:30432"
|
|
config.memory.model = "hermes-agent"
|
|
config.memory.api_key_env = "TEST_HERMES_KEY"
|
|
captured: dict = {}
|
|
|
|
async def handler(request: httpx.Request) -> httpx.Response:
|
|
captured["url"] = str(request.url)
|
|
captured["authorization"] = request.headers.get("Authorization")
|
|
captured["payload"] = json.loads(request.content)
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"choices": [
|
|
{"message": {"role": "assistant", "content": "[]"}}
|
|
]
|
|
},
|
|
)
|
|
|
|
client = HermesClient(config)
|
|
await client._client.aclose()
|
|
client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
|
try:
|
|
with patch.dict(os.environ, {"TEST_HERMES_KEY": "test-secret"}):
|
|
response = await client.chat(
|
|
[{"role": "user", "content": "Extract memories"}],
|
|
temperature=0.0,
|
|
)
|
|
finally:
|
|
await client.aclose()
|
|
|
|
self.assertEqual(captured["url"], "http://hermes.test:30432/v1/chat/completions")
|
|
self.assertEqual(captured["authorization"], "Bearer test-secret")
|
|
self.assertEqual(captured["payload"]["model"], "hermes-agent")
|
|
self.assertFalse(captured["payload"]["stream"])
|
|
self.assertEqual(response.content, "[]")
|
|
|
|
async def test_list_models_is_authenticated(self) -> None:
|
|
config = Config()
|
|
config.memory.provider = "hermes"
|
|
config.memory.base_url = "http://hermes.test:30432"
|
|
config.memory.api_key_env = "TEST_HERMES_KEY"
|
|
authorization = ""
|
|
|
|
async def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal authorization
|
|
authorization = request.headers.get("Authorization", "")
|
|
return httpx.Response(200, json={"data": [{"id": "hermes-agent"}]})
|
|
|
|
client = HermesClient(config)
|
|
await client._client.aclose()
|
|
client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
|
try:
|
|
with patch.dict(os.environ, {"TEST_HERMES_KEY": "test-secret"}):
|
|
models = await client.list_models()
|
|
finally:
|
|
await client.aclose()
|
|
|
|
self.assertEqual(authorization, "Bearer test-secret")
|
|
self.assertEqual(models, ["hermes-agent"])
|
|
|
|
|
|
class FakeMemoryClient:
|
|
async def chat(self, messages: list[dict], **kwargs) -> ChatResponse:
|
|
return ChatResponse(content=json.dumps([
|
|
{
|
|
"key": "music_preference",
|
|
"value": "Florian prefers jazz music",
|
|
"category": "preferences",
|
|
}
|
|
]))
|
|
|
|
|
|
class HermesMemoryMiddlewareTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_short_explicit_request_is_extracted_and_persisted(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_directory:
|
|
config = Config(paths=PathsConfig(base_dir=Path(temporary_directory)))
|
|
config.memory.provider = "hermes"
|
|
config.memory.model = "hermes-agent"
|
|
middleware = MemoryMiddleware(FakeMemoryClient(), config)
|
|
context = TurnContext(turn_id="turn-1", speaker_name="Florian")
|
|
|
|
await middleware.on_finish(
|
|
context,
|
|
"Merke dir: Ich mag Jazz.",
|
|
"Das merke ich mir.",
|
|
)
|
|
|
|
self.assertEqual(len(context._deferred), 1)
|
|
await context._deferred[0]
|
|
personal_file = Path(temporary_directory) / "memory_florian.json"
|
|
stored = json.loads(personal_file.read_text(encoding="utf-8"))
|
|
self.assertEqual(
|
|
stored["entries"]["music_preference"]["value"],
|
|
"Florian prefers jazz music",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |