Add results to completions

Author: MathyouMB

requirements.txt

@@ -32,11 +32,13 @@ prompt_toolkit==3.0.51
 psycopg2-binary==2.9.10
 PyMuPDF==1.25.5
 pytest==8.3.5
+pytest-describe==2.2.0
 python-dateutil==2.9.0.post0
 PyYAML==6.0.2
 redis==5.2.1
 regex==2024.11.6
 requests==2.32.3
+result==0.17.0
 safetensors==0.5.3
 scikit-learn==1.6.1
 scipy==1.15.2

src/apps/dev_cli.py

@@ -5,6 +5,7 @@
 from completions import CompletionService
 from infrastructure.db.source_repository import SourceRepository
 from config.logger import logger
+from result import Ok
 
 if __name__ == "__main__":
     with SourceRepository() as repository:
@@ -14,5 +15,12 @@
             q = input("\nAsk a question (or 'exit'): ")
             if q.lower() in {"exit", "quit"}:
                 break
-            answer = service.create(q)
-            logger.info("dev_cli.answer_displayed", answer=answer)
+
+            result = service.create(q)
+
+            if isinstance(result, Ok):
+                logger.info("dev_cli.answer_displayed", answer=result.ok_value)
+                print(f"\nAnswer: {result.ok_value}")
+            else:
+                logger.error("dev_cli.answer_failed", error=repr(result.err_value))
+                print(f"\nError: {result.err_value}")

src/completions/services/completion_service.py

@@ -1,3 +1,4 @@
+from result import Ok, Result
 from completions.prompts import create_prompt
 from infrastructure.ollama import create_completion
 from infrastructure.db.source_repository import SourceRepository
@@ -9,16 +10,26 @@ class CompletionService:
     def __init__(self, repository: SourceRepository):
         self.repository = repository
 
-    def create(self, query, k=3):
+    def create(self, query: str, k: int = 3) -> Result[str, Exception]:
         logger.info("completion_service.create", query=query)
 
         prompt_context = ContextGenerationService(query, self.repository).process(k)
         logger.info("completion_service.context_generated", query=query)
 
-        answer = create_completion(create_prompt(prompt_context, query))
-        logger.info(
-            "completion_service.answer_generated",
-            query=query,
-            answer=answer[0:10] + "...",
-        )
-        return answer
+        prompt = create_prompt(query, prompt_context)
+        result = create_completion(prompt)
+
+        if isinstance(result, Ok):
+            logger.info(
+                "completion_service.answer_generated",
+                query=query,
+                answer=result.ok_value[:10] + "...",
+            )
+        else:
+            logger.error(
+                "completion_service.answer_generation_error",
+                query=query,
+                error=result.err_value,
+            )
+
+        return result

src/config/logger.py

@@ -25,7 +25,6 @@ def __init__(self, name: Optional[str] = None):
         self.logger = logging.getLogger(name or __name__)
         self.logger.propagate = False  # 🔒 Prevent double logging
 
-        # Only add handler if none exist
         if not self.logger.handlers:
             handler = logging.StreamHandler()
             formatter = ColorFormatter("%(asctime)s [%(levelname)s] %(message)s")
@@ -48,7 +47,7 @@ def debug(self, event: str, **data: Any):
 
     def _format(self, event: str, data: Dict[str, Any]) -> str:
         structured = {"event": event, **data}
-        return json.dumps(structured)
+        return json.dumps(structured, default=str)
 
 
 logger = SemanticLogger(__name__)

src/infrastructure/ollama.py

@@ -1,10 +1,16 @@
 import requests
+from result import Ok, Err, Result
 from config import OLLAMA_MODEL, OLLAMA_URL
 
 
-def create_completion(prompt):
-    res = requests.post(
-        f"{OLLAMA_URL}/api/generate",
-        json={"model": OLLAMA_MODEL, "prompt": prompt, "stream": False},
-    )
-    return res.json()["response"]
+def create_completion(prompt: str) -> Result[str, Exception]:
+    try:
+        res = requests.post(
+            f"{OLLAMA_URL}/api/generate",
+            json={"model": OLLAMA_MODEL, "prompt": prompt, "stream": False},
+            timeout=30,
+        )
+        res.raise_for_status()
+        return Ok(res.json()["response"])
+    except (requests.RequestException, ValueError, KeyError) as e:
+        return Err(e)

tests/completions/services/test_completion_service.py

@@ -1,5 +1,6 @@
 import pytest
 from unittest.mock import MagicMock, patch
+from result import Ok
 from completions import CompletionService
 
 
@@ -27,23 +28,24 @@ def service(fake_repository):
 
 @patch(
     "completions.services.completion_service.create_completion",
-    return_value="LLM Answer",
+    return_value=Ok("LLM Answer"),
 )
 @patch(
     "completions.services.completion_service.create_prompt",
     return_value="Generated Prompt",
 )
-def test_create(mock_create_prompt, mock_create_completion, service, capsys):
-    """Should call the LLM with the correct prompt and print the answer."""
+def test_create(mock_create_prompt, mock_create_completion, service):
+    """Should call the LLM with the correct prompt and return Ok result."""
 
     result = service.create("What is AI?", k=3)
 
     service.repository.get_top_k_chunks_by_similarity.assert_called_once_with(
         "What is AI?", 3
     )
     mock_create_prompt.assert_called_once_with(
-        "Chunk 1\n\nChunk 2\n\nChunk 3", "What is AI?"
+        "What is AI?", "Chunk 1\n\nChunk 2\n\nChunk 3"
     )
     mock_create_completion.assert_called_once_with("Generated Prompt")
 
-    assert result == "LLM Answer"
+    assert isinstance(result, Ok)
+    assert result.ok_value == "LLM Answer"