Create workflow.py

workflow.py

@@ -0,0 +1,214 @@
+# workflow.py
+
+import time
+from datetime import datetime
+from typing import Dict
+
+from langchain_core.messages import AIMessage, HumanMessage
+from langgraph.graph import END, StateGraph
+from langgraph.graph.message import add_messages
+
+from processor import EnhancedCognitiveProcessor
+from config import ResearchConfig
+
+import logging
+logger = logging.getLogger(__name__)
+
+class ResearchWorkflow:
+    """
+    Defines a multi-step research workflow using a state graph.
+    """
+    def __init__(self) -> None:
+        self.processor = EnhancedCognitiveProcessor()
+        self.workflow = StateGraph()
+        self._build_workflow()
+        self.app = self.workflow.compile()
+
+    def _build_workflow(self) -> None:
+        self.workflow.add_node("ingest", self.ingest_query)
+        self.workflow.add_node("retrieve", self.retrieve_documents)
+        self.workflow.add_node("analyze", self.analyze_content)
+        self.workflow.add_node("validate", self.validate_output)
+        self.workflow.add_node("refine", self.refine_results)
+        self.workflow.set_entry_point("ingest")
+        self.workflow.add_edge("ingest", "retrieve")
+        self.workflow.add_edge("retrieve", "analyze")
+        self.workflow.add_conditional_edges(
+            "analyze",
+            self._quality_check,
+            {"valid": "validate", "invalid": "refine"}
+        )
+        self.workflow.add_edge("validate", END)
+        self.workflow.add_edge("refine", "retrieve")
+        # Extended node for multi-modal enhancement
+        self.workflow.add_node("enhance", self.enhance_analysis)
+        self.workflow.add_edge("validate", "enhance")
+        self.workflow.add_edge("enhance", END)
+
+    def ingest_query(self, state: Dict) -> Dict:
+        try:
+            query = state["messages"][-1].content
+            # Retrieve the domain from the state's context (defaulting to Biomedical Research)
+            domain = state.get("context", {}).get("domain", "Biomedical Research")
+            new_context = {"raw_query": query, "domain": domain, "refine_count": 0, "refinement_history": []}
+            logger.info(f"Query ingested. Domain: {domain}")
+            return {
+                "messages": [AIMessage(content="Query ingested successfully")],
+                "context": new_context,
+                "metadata": {"timestamp": datetime.now().isoformat()}
+            }
+        except Exception as e:
+            logger.exception("Error during query ingestion.")
+            return self._error_state(f"Ingestion Error: {str(e)}")
+
+    def retrieve_documents(self, state: Dict) -> Dict:
+        try:
+            query = state["context"]["raw_query"]
+            # For demonstration, we use an empty document list.
+            # Replace this with actual retrieval logic as needed.
+            docs = []
+            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(),
+                    "refine_count": state["context"].get("refine_count", 0),
+                    "refinement_history": state["context"].get("refinement_history", []),
+                    "domain": state["context"].get("domain", "Biomedical Research")
+                }
+            }
+        except Exception as e:
+            logger.exception("Error during document retrieval.")
+            return self._error_state(f"Retrieval Error: {str(e)}")
+
+    def analyze_content(self, state: Dict) -> Dict:
+        try:
+            domain = state["context"].get("domain", "Biomedical Research").strip().lower()
+            fallback_analyses = ResearchConfig.DOMAIN_FALLBACKS
+            if domain in fallback_analyses:
+                logger.info(f"Using fallback analysis for domain: {state['context'].get('domain')}")
+                return {
+                    "messages": [AIMessage(content=fallback_analyses[domain].strip())],
+                    "context": state["context"]
+                }
+            else:
+                docs = state["context"].get("documents", [])
+                docs_text = "\n\n".join([d.page_content for d in docs])
+                domain_prompt = ResearchConfig.DOMAIN_PROMPTS.get(domain, "")
+                full_prompt = f"{domain_prompt}\n\n" + ResearchConfig.ANALYSIS_TEMPLATE.format(context=docs_text)
+                response = self.processor.process_query(full_prompt)
+                if "error" in response:
+                    logger.error("Backend response error during analysis.")
+                    return self._error_state(response["error"])
+                logger.info("Content analysis completed.")
+                return {
+                    "messages": [AIMessage(content=response.get('choices', [{}])[0].get('message', {}).get('content', ''))],
+                    "context": state["context"]
+                }
+        except Exception as e:
+            logger.exception("Error during content analysis.")
+            return self._error_state(f"Analysis Error: {str(e)}")
+
+    def validate_output(self, state: Dict) -> Dict:
+        try:
+            analysis = state["messages"][-1].content
+            validation_prompt = (
+                f"Validate the following research analysis:\n{analysis}\n\n"
+                "Check for:\n"
+                "1. Technical accuracy\n"
+                "2. Citation support (are claims backed by evidence?)\n"
+                "3. Logical consistency\n"
+                "4. Methodological soundness\n\n"
+                "Respond with 'VALID: [brief justification]' or 'INVALID: [brief justification]'."
+            )
+            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', '')}")]
+            }
+        except Exception as e:
+            logger.exception("Error during output validation.")
+            return self._error_state(f"Validation Error: {str(e)}")
+
+    def refine_results(self, state: Dict) -> Dict:
+        try:
+            current_count = state["context"].get("refine_count", 0)
+            state["context"]["refine_count"] = current_count + 1
+            refinement_history = state["context"].setdefault("refinement_history", [])
+            current_analysis = state["messages"][-1].content
+            refinement_history.append(current_analysis)
+            difficulty_level = max(0, 3 - state["context"]["refine_count"])
+            logger.info(f"Refinement iteration: {state['context']['refine_count']}, Difficulty level: {difficulty_level}")
+
+            if state["context"]["refine_count"] >= 3:
+                meta_prompt = (
+                    "You are given the following series of refinement outputs:\n" +
+                    "\n---\n".join(refinement_history) +
+                    "\n\nSynthesize the above into a final, concise, and high-quality technical analysis report. "
+                    "Focus on the key findings and improvements made across the iterations. Do not introduce new ideas; just synthesize the improvements. Ensure the report is well-structured and easy to understand."
+                )
+                meta_response = self.processor.process_query(meta_prompt)
+                logger.info("Meta-refinement completed.")
+                return {
+                    "messages": [AIMessage(content=meta_response.get('choices', [{}])[0].get('message', {}).get('content', ''))],
+                    "context": state["context"]
+                }
+            else:
+                refinement_prompt = (
+                    f"Refine this analysis (current difficulty level: {difficulty_level}):\n{current_analysis}\n\n"
+                    "First, critically evaluate the analysis and identify its weaknesses, such as inaccuracies, unsupported claims, or lack of clarity. Summarize these weaknesses in a short paragraph.\n\n"
+                    "Then, improve the following aspects:\n"
+                    "1. Technical precision\n"
+                    "2. Empirical grounding\n"
+                    "3. Theoretical coherence\n\n"
+                    "Use a structured difficulty gradient approach (similar to LADDER) to produce a simpler yet more accurate variant, addressing the weaknesses identified."
+                )
+                response = self.processor.process_query(refinement_prompt)
+                logger.info("Refinement completed.")
+                return {
+                    "messages": [AIMessage(content=response.get('choices', [{}])[0].get('message', {}).get('content', ''))],
+                    "context": state["context"]
+                }
+        except Exception as e:
+            logger.exception("Error during refinement.")
+            return self._error_state(f"Refinement Error: {str(e)}")
+
+    def _quality_check(self, state: Dict) -> str:
+        refine_count = state["context"].get("refine_count", 0)
+        if refine_count >= 3:
+            logger.warning("Refinement limit reached. Forcing valid outcome.")
+            return "valid"
+        content = state["messages"][-1].content
+        quality = "valid" if "VALID" in content else "invalid"
+        logger.info(f"Quality check returned: {quality}")
+        return quality
+
+    def _error_state(self, message: str) -> Dict:
+        logger.error(message)
+        return {
+            "messages": [{"content": f"❌ {message}"}],
+            "context": {"error": True},
+            "metadata": {"status": "error"}
+        }
+
+    def enhance_analysis(self, state: Dict) -> Dict:
+        try:
+            analysis = state["messages"][-1].content
+            enhanced = f"{analysis}\n\n## Multi-Modal Insights\n"
+            if "images" in state["context"]:
+                enhanced += "### Visual Evidence\n"
+                for img in state["context"]["images"]:
+                    enhanced += f"![Relevant visual]({img})\n"
+            if "code" in state["context"]:
+                enhanced += "### Code Artifacts\n```python\n"
+                for code in state["context"]["code"]:
+                    enhanced += f"{code}\n"
+                enhanced += "```"
+            return {
+                "messages": [{"content": enhanced}],
+                "context": state["context"]
+            }
+        except Exception as e:
+            logger.exception("Error during multi-modal enhancement.")
+            return self._error_state(f"Enhancement Error: {str(e)}")