Improvements

src/llm/processor.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from typing import Any, Dict, Optional
 
@@ -7,6 +8,8 @@ from pydantic import ValidationError
 from .models import ExtractionResult
 from .prompts import EXTRACTION_SYSTEM_PROMPT, NOISE_FILTER_SYSTEM_PROMPT
 
+logger = logging.getLogger(__name__)
+
 
 class LLMProcessor:
     def __init__(
@@ -47,7 +50,7 @@ class LLMProcessor:
                 # but we can ensure the client is initialized.
                 pass
         except Exception as e:
-            print(f"Error initializing LLM client for backend {backend}: {e}")
+            logger.error(f"Error initializing LLM client for backend {backend}: {e}")
             raise
 
         self.model = model or os.environ.get("LLM_MODEL", "gpt-4o")
@@ -56,73 +59,98 @@ class LLMProcessor:
         self,
         system_prompt: str,
         user_prompt: str,
+        context: Optional[str] = None,
         response_format: Optional[Any] = None,
     ) -> str:
         """
         Generic method to call the LLM.
         """
+        messages = [
+            {"role": "system", "content": system_prompt},
+        ]
+        if context:
+            messages.append(
+                {
+                    "role": "system",
+                    "content": f"Context from previous conversation:\n{context}",
+                }
+            )
+
+        messages.append({"role": "user", "content": user_prompt})
+
         try:
             response = self.client.chat.completions.create(
                 model=self.model,
-                messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt},
-                ],
+                messages=messages,
                 response_format=response_format,
                 extra_body={"include_reasoning": False},
             )
             return response.choices[0].message.content
         except Exception as e:
-            print(f"LLM Error: {e}")
+            logger.error(f"LLM Error: {e}")
             return ""
 
-    def filter_transcript(self, text: str) -> str:
+    def filter_transcript(self, text: str, context: Optional[str] = None) -> str:
         """
         Stage 1: Raw Transcript -> Filtered Text.
         """
-        result = self._call_llm(NOISE_FILTER_SYSTEM_PROMPT, text)
-        print(f"LLM Processor (Filter): {text} -> {result}")
+        result = self._call_llm(NOISE_FILTER_SYSTEM_PROMPT, text, context=context)
+        logger.info(f"LLM Processor (Filter): {text} -> {result}")
         return result
 
-    def extract_structured_data(self, filtered_text: str) -> ExtractionResult:
+    def extract_structured_data(
+        self, filtered_text: str, context: Optional[str] = None
+    ) -> ExtractionResult:
         """
         Stage 2: Filtered Text -> Structured Data.
         """
-        print(f"LLM Processor (Extract): Calling extraction for: {filtered_text}")
+        logger.info(f"LLM Processor (Extract): Calling extraction for: {filtered_text}")
         try:
             # Using standard chat.completions.create with JSON mode for better compatibility with vLLM
-            print("LLM Processor (Extract): Sending request to backend...")
+            logger.info("LLM Processor (Extract): Sending request to backend...")
+
+            messages = [
+                {"role": "system", "content": EXTRACTION_SYSTEM_PROMPT},
+            ]
+            if context:
+                messages.append(
+                    {
+                        "role": "system",
+                        "content": f"Context from previous conversation:\n{context}",
+                    }
+                )
+            messages.append({"role": "user", "content": filtered_text})
+
             response = self.client.chat.completions.create(
                 model=self.model,
-                messages=[
-                    {"role": "system", "content": EXTRACTION_SYSTEM_PROMPT},
-                    {"role": "user", "content": filtered_text},
-                ],
+                messages=messages,
                 response_format={"type": "json_object"},
                 extra_body={"include_reasoning": False},
             )
-            print("LLM Processor (Extract): Response received from backend.")
+            logger.info("LLM Processor (Extract): Response received from backend.")
 
             import json
 
             content = response.choices[0].message.content
-            print(f"LLM Processor (Extract): Raw JSON response: {content}")
+            logger.info(f"LLM Processor (Extract): Raw JSON response: {content}")
             data = json.loads(content)
 
             # Map the JSON data to the Pydantic model
             return ExtractionResult(**data)
 
         except Exception as e:
-            print(f"Extraction Error: {e}")
+            logger.error(f"Extraction Error: {e}")
             # Return an empty ExtractionResult if parsing fails
             return ExtractionResult()
 
-    def process_pipeline(self, raw_text: str) -> ExtractionResult:
+    def process_pipeline(
+        self, raw_text: str, context: Optional[str] = None
+    ) -> ExtractionResult:
         """
         Executes the two-stage pipeline: Raw Transcript -> Filtered Text -> Structured Data.
         """
-        filtered_text = self.filter_transcript(raw_text)
+        filtered_text = self.filter_transcript(raw_text, context=context)
         if not filtered_text:
             return ExtractionResult()
 
-        return self.extract_structured_data(filtered_text)
+        return self.extract_structured_data(filtered_text, context=context)