Update llm_notes_generator.py

notes/llm_notes_generator.py

@@ -34,8 +34,8 @@ logger = logging.getLogger(__name__)
 
 class Settings(BaseSettings):
 	"""Application settings loaded from environment variables or .env file."""
-	mistral_api_key: str = os.getenv('MISTRAL_API_KEY', '')
-	api_url: str = "https://api.mistral.ai/v1/chat/completions"
+	openrouter_api_key: str = os.getenv('OPENROUTER_API_KEY', '')
+	api_url: str = "https://openrouter.ai/api/v1/chat/completions"
 	output_dir: str = "data/generated_notes"
 	trial_balance_json: str = "data/output1/parsed_trial_balance.json"
 
@@ -61,20 +61,22 @@ class GeneratedNote(BaseModel):
 
 class FlexibleFinancialNoteGenerator:
     def __init__(self):
-        self.mistral_api_key = settings.mistral_api_key
-        if not self.mistral_api_key:
-            logger.error("MISTRAL_API_KEY not found in .env file")
-            raise ValueError("MISTRAL_API_KEY not found in .env file")
+        self.openrouter_api_key = settings.openrouter_api_key
+        if not self.openrouter_api_key:
+            logger.error("OPENROUTER_API_KEY not found in .env file")
+            raise ValueError("OPENROUTER_API_KEY not found in .env file")
         self.api_url = settings.api_url
         self.headers = {
-            "Authorization": f"Bearer {self.mistral_api_key}",
-            "Content-Type": "application/json"
+            "Authorization": f"Bearer {self.openrouter_api_key}",
+            "Content-Type": "application/json",
+            "HTTP-Referer": "https://localhost:3000",
+            "X-Title": "Financial Note Generator"
         }
         self.note_templates = self.load_note_templates()
         self.account_patterns = self._init_account_patterns()
         self.recommended_models = [
-            "mistral-large-latest",
-            "mistral-medium-latest"
+            "mistralai/mixtral-8x7b-instruct",
+            "mistralai/mistral-7b-instruct-v0.2"
         ]
 
     def _init_account_patterns(self) -> Dict[str, Dict[str, Any]]:
@@ -338,8 +340,8 @@ class FlexibleFinancialNoteGenerator:
         
         return prompt
     
-    def call_mistral_api(self, prompt: str) -> Optional[str]:
-        """Make API call to Mistral with model fallback"""
+    def call_openrouter_api(self, prompt: str) -> Optional[str]:
+        """Make API call to OpenRouter with model fallback"""
         for model in self.recommended_models:
             logger.info(f"Trying model: {model}")
             payload = {
@@ -357,7 +359,7 @@ class FlexibleFinancialNoteGenerator:
                     self.api_url,
                     headers=self.headers,
                     json=payload,
-                    timeout=30
+                    timeout=30  # <-- Add timeout here!
                 )
                 response.raise_for_status()
                 result = response.json()
@@ -448,7 +450,7 @@ class FlexibleFinancialNoteGenerator:
             logger.error("Failed to build prompt")
             return False
         
-        response = self.call_mistral_api(prompt)
+        response = self.call_openrouter_api(prompt)
         if not response:
             logger.error("Failed to get API response")
             return False
@@ -473,7 +475,7 @@ class FlexibleFinancialNoteGenerator:
             if not prompt:
                 results[note_number] = False
                 continue
-            response = self.call_mistral_api(prompt)
+            response = self.call_openrouter_api(prompt)
             if not response:
                 results[note_number] = False
                 continue
@@ -586,4 +588,4 @@ def main() -> None:
         sys.exit(1)
 
 if __name__ == "__main__":
-    main()
+    main()