Update app.py

app.py

@@ -1,5 +1,5 @@
 # ------------------------------
-# Enhanced NeuroResearch AI System
+# Enhanced NeuroResearch AI System with Refinement Counter
 # ------------------------------
 import logging
 import os
@@ -7,10 +7,12 @@ import re
 import hashlib
 import json
 import time
+import sys
 from datetime import datetime
+import chromdb
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import List, Dict, Any, Optional, Sequence
-import chromadb
+
 import requests
 import streamlit as st
 
@@ -25,11 +27,14 @@ from langgraph.graph.message import add_messages
 from typing_extensions import TypedDict, Annotated
 from langchain.tools.retriever import create_retriever_tool
 
+# Optionally increase the recursion limit if needed
+sys.setrecursionlimit(100)
+
 # ------------------------------
 # Logging Configuration
 # ------------------------------
 logging.basicConfig(
-    level=logging.INFO, 
+    level=logging.INFO,
     format="%(asctime)s [%(levelname)s] %(message)s"
 )
 logger = logging.getLogger(__name__)
@@ -53,11 +58,11 @@ class ResearchConfig:
     MAX_CONCURRENT_REQUESTS = 5
     EMBEDDING_DIMENSIONS = 1536
     DOCUMENT_MAP = {
-        "Research Report: Results of a New AI Model Improving Image Recognition Accuracy to 98%": 
+        "Research Report: Results of a New AI Model Improving Image Recognition Accuracy to 98%":
             "CV-Transformer Hybrid Architecture",
-        "Academic Paper Summary: Why Transformers Became the Mainstream Architecture in Natural Language Processing": 
+        "Academic Paper Summary: Why Transformers Became the Mainstream Architecture in Natural Language Processing":
             "Transformer Architecture Analysis",
-        "Latest Trends in Machine Learning Methods Using Quantum Computing": 
+        "Latest Trends in Machine Learning Methods Using Quantum Computing":
             "Quantum ML Frontiers"
     }
     ANALYSIS_TEMPLATE = (
@@ -85,7 +90,7 @@ if not ResearchConfig.DEEPSEEK_API_KEY:
 # ------------------------------
 class QuantumDocumentManager:
     """
-    Manages the creation of Chroma collections from raw document texts.
+    Manages creation of Chroma collections from raw document texts.
     """
     def __init__(self) -> None:
         try:
@@ -197,7 +202,7 @@ class CognitiveProcessor:
 
     def process_query(self, prompt: str) -> Dict:
         """
-        Process a query by sending multiple API requests in parallel.
+        Processes a query by sending multiple API requests in parallel.
         """
         futures = []
         for _ in range(3):  # Triple redundancy for reliability
@@ -254,7 +259,6 @@ class CognitiveProcessor:
         if not valid_results:
             logger.error("All API requests failed.")
             return {"error": "All API requests failed"}
-        # Choose the response with the longest content
         return max(valid_results, key=lambda x: len(x.get('choices', [{}])[0].get('message', {}).get('content', '')))
 
 # ------------------------------
@@ -291,14 +295,16 @@ class ResearchWorkflow:
 
     def ingest_query(self, state: AgentState) -> Dict:
         """
-        Ingests the research query.
+        Ingests the research query and initializes the refinement counter.
         """
         try:
             query = state["messages"][-1].content
+            # Initialize context with raw query and refinement counter
+            new_context = {"raw_query": query, "refine_count": 0}
             logger.info("Query ingested.")
             return {
                 "messages": [AIMessage(content="Query ingested successfully")],
-                "context": {"raw_query": query},
+                "context": new_context,
                 "metadata": {"timestamp": datetime.now().isoformat()}
             }
         except Exception as e:
@@ -314,7 +320,7 @@ class ResearchWorkflow:
             logger.info(f"Retrieved {len(docs)} documents for query.")
             return {
                 "messages": [AIMessage(content=f"Retrieved {len(docs)} documents")],
-                "context": {"documents": docs, "retrieval_time": time.time()}
+                "context": {"documents": docs, "retrieval_time": time.time(), "refine_count": state["context"].get("refine_count", 0)}
             }
         except Exception as e:
             return self._error_state(f"Retrieval Error: {str(e)}")
@@ -333,7 +339,7 @@ class ResearchWorkflow:
             logger.info("Content analysis completed.")
             return {
                 "messages": [AIMessage(content=response.get('choices', [{}])[0].get('message', {}).get('content', ''))],
-                "context": {"analysis": response}
+                "context": {"analysis": response, "refine_count": state["context"].get("refine_count", 0)}
             }
         except Exception as e:
             return self._error_state(f"Analysis Error: {str(e)}")
@@ -351,18 +357,18 @@ class ResearchWorkflow:
         response = self.processor.process_query(validation_prompt)
         logger.info("Output validation completed.")
         return {
-            "messages": [
-                AIMessage(
-                    content=analysis +
-                    f"\n\nValidation: {response.get('choices', [{}])[0].get('message', {}).get('content', '')}"
-                )
-            ]
+            "messages": [AIMessage(content=analysis + f"\n\nValidation: {response.get('choices', [{}])[0].get('message', {}).get('content', '')}")]
         }
 
     def refine_results(self, state: AgentState) -> Dict:
         """
         Refines the analysis report if validation fails.
+        Increments the refinement counter to limit infinite loops.
         """
+        current_count = state["context"].get("refine_count", 0)
+        state["context"]["refine_count"] = current_count + 1
+        logger.info(f"Refinement iteration: {state['context']['refine_count']}")
+
         refinement_prompt = (
             f"Refine this analysis:\n{state['messages'][-1].content}\n\n"
             "Improve:\n1. Technical precision\n2. Empirical grounding\n3. Theoretical coherence"
@@ -370,18 +376,19 @@ class ResearchWorkflow:
         response = self.processor.process_query(refinement_prompt)
         logger.info("Refinement completed.")
         return {
-            "messages": [
-                AIMessage(
-                    content=response.get('choices', [{}])[0].get('message', {}).get('content', '')
-                )
-            ],
+            "messages": [AIMessage(content=response.get('choices', [{}])[0].get('message', {}).get('content', ''))],
             "context": state["context"]
         }
 
     def _quality_check(self, state: AgentState) -> str:
         """
         Checks whether the analysis report is valid.
+        Forces a valid state if the refinement count exceeds a threshold.
         """
+        refine_count = state["context"].get("refine_count", 0)
+        if refine_count >= 3:
+            logger.warning("Refinement limit reached. Forcing valid outcome to prevent infinite recursion.")
+            return "valid"
         content = state["messages"][-1].content
         quality = "valid" if "VALID" in content else "invalid"
         logger.info(f"Quality check returned: {quality}")