""" LangGraph Workflow for SPARKNET Implements cyclic multi-agent workflows with StateGraph """ from typing import Literal, Dict, Any, Optional from datetime import datetime from loguru import logger from langgraph.graph import StateGraph, END from langgraph.checkpoint.memory import MemorySaver from langchain_core.messages import HumanMessage, AIMessage, SystemMessage from .langgraph_state import ( AgentState, ScenarioType, TaskStatus, WorkflowOutput, create_initial_state, state_to_output, ) from ..llm.langchain_ollama_client import LangChainOllamaClient class SparknetWorkflow: """ LangGraph-powered workflow orchestrator for SPARKNET. Implements cyclic workflow with conditional routing: START → PLANNER → ROUTER → [scenario executors] → CRITIC ↑ ↓ └────────── REFINE ←──────────────────────┘ """ def __init__( self, llm_client: LangChainOllamaClient, planner_agent: Optional[Any] = None, critic_agent: Optional[Any] = None, memory_agent: Optional[Any] = None, vision_ocr_agent: Optional[Any] = None, quality_threshold: float = 0.85, max_iterations: int = 3, ): self.llm_client = llm_client self.planner_agent = planner_agent self.critic_agent = critic_agent self.memory_agent = memory_agent self.vision_ocr_agent = vision_ocr_agent self.quality_threshold = quality_threshold self.max_iterations = max_iterations self.graph = self._build_graph() self.checkpointer = MemorySaver() self.app = self.graph.compile(checkpointer=self.checkpointer) if vision_ocr_agent: logger.info("Initialized SparknetWorkflow with LangGraph StateGraph and VisionOCR support") else: logger.info("Initialized SparknetWorkflow with LangGraph StateGraph") def _build_graph(self) -> StateGraph: workflow = StateGraph(AgentState) workflow.add_node("planner", self._planner_node) workflow.add_node("router", self._router_node) workflow.add_node("executor", self._executor_node) workflow.add_node("critic", self._critic_node) workflow.add_node("refine", self._refine_node) workflow.add_node("finish", self._finish_node) workflow.set_entry_point("planner") workflow.add_edge("planner", "router") workflow.add_edge("router", "executor") workflow.add_edge("executor", "critic") workflow.add_conditional_edges( "critic", self._should_refine, { "refine": "refine", "finish": "finish", } ) workflow.add_edge("refine", "planner") workflow.add_edge("finish", END) return workflow async def _planner_node(self, state: AgentState) -> AgentState: logger.info(f"PLANNER node processing task: {state['task_id']}") state["status"] = TaskStatus.PLANNING state["current_agent"] = "PlannerAgent" # Retrieve relevant context from memory context_docs = [] if self.memory_agent: try: logger.info("Retrieving relevant context from memory...") context_docs = await self.memory_agent.retrieve_relevant_context( query=state["task_description"], context_type="all", top_k=3, scenario_filter=state["scenario"], min_quality_score=0.8 ) if context_docs: logger.info(f"Retrieved {len(context_docs)} relevant memories") # Add context to state for reference state["agent_outputs"]["memory_context"] = [ {"content": doc.page_content, "metadata": doc.metadata} for doc in context_docs ] except Exception as e: logger.warning(f"Memory retrieval failed: {e}") system_msg = SystemMessage(content="Decompose the task into executable subtasks.") # Add memory context to user message if available context_text = "" if context_docs: context_text = "\n\nRelevant past experiences:\n" for i, doc in enumerate(context_docs, 1): context_text += f"\n{i}. {doc.page_content[:200]}..." user_msg = HumanMessage( content=f"Task: {state['task_description']}\nScenario: {state['scenario']}{context_text}" ) llm = self.llm_client.get_llm(complexity="complex") if self.planner_agent: from ..agents.base_agent import Task task = Task( id=state["task_id"], description=state["task_description"], metadata={"scenario": state["scenario"].value} ) result_task = await self.planner_agent.process_task(task) if result_task.status == "completed": state["subtasks"] = [ { "id": st.id, "description": st.description, "agent_type": st.agent_type, "dependencies": st.dependencies, } for st in result_task.result["task_graph"].subtasks.values() ] state["execution_order"] = result_task.result["execution_order"] response_msg = AIMessage(content=f"Created plan with {len(state['subtasks'])} subtasks") state["messages"].append(response_msg) else: response = await llm.ainvoke([system_msg, user_msg]) state["messages"].append(response) state["subtasks"] = [ {"id": "subtask_1", "description": "Execute primary task", "agent_type": "ExecutorAgent", "dependencies": []} ] state["execution_order"] = [["subtask_1"]] logger.info(f"Planning completed: {len(state.get('subtasks', []))} subtasks created") return state async def _router_node(self, state: AgentState) -> AgentState: logger.info(f"ROUTER node routing for scenario: {state['scenario']}") state["current_agent"] = "Router" scenario = state["scenario"] routing_msg = AIMessage(content=f"Routing to {scenario.value} workflow agents") state["messages"].append(routing_msg) state["agent_outputs"]["router"] = { "scenario": scenario.value, "agents_to_use": self._get_scenario_agents(scenario) } return state async def _executor_node(self, state: AgentState) -> AgentState: logger.info(f"EXECUTOR node executing for scenario: {state['scenario']}") state["status"] = TaskStatus.EXECUTING state["current_agent"] = "Executor" scenario = state["scenario"] # Route to scenario-specific pipeline if scenario == ScenarioType.PATENT_WAKEUP: logger.info("🎯 Routing to Patent Wake-Up pipeline") return await self._execute_patent_wakeup(state) # Generic execution for other scenarios agents = self._get_scenario_agents(scenario) # Get scenario-specific tools from ..tools.langchain_tools import get_vista_tools tools = get_vista_tools(scenario.value) logger.info(f"Loaded {len(tools)} tools for scenario: {scenario.value}") # Bind tools to LLM llm = self.llm_client.get_llm(complexity="standard") llm_with_tools = llm.bind_tools(tools) # Build execution prompt with tool information tool_descriptions = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools]) execution_prompt = HumanMessage( content=f"""Execute the following task using the available tools when needed: Task: {state['task_description']} Scenario: {scenario.value} Available tools: {tool_descriptions} Provide detailed results.""" ) # Execute with tool support response = await llm_with_tools.ainvoke([execution_prompt]) state["messages"].append(response) # Check if tools were called tool_calls = [] if hasattr(response, 'tool_calls') and response.tool_calls: logger.info(f"LLM requested {len(response.tool_calls)} tool calls") for tool_call in response.tool_calls: tool_name = tool_call.get('name', 'unknown') tool_calls.append(tool_name) logger.info(f"Tool called: {tool_name}") state["agent_outputs"]["executor"] = { "result": response.content, "agents_used": agents, "tools_available": [tool.name for tool in tools], "tools_called": tool_calls, } state["final_output"] = response.content logger.info("Execution completed") return state async def _execute_patent_wakeup(self, state: AgentState) -> AgentState: """ Execute Patent Wake-Up scenario pipeline. Sequential execution: Document → Market → Matchmaking → Outreach """ logger.info("🚀 Executing Patent Wake-Up pipeline") # Import scenario1 agents from ..agents.scenario1 import ( DocumentAnalysisAgent, MarketAnalysisAgent, MatchmakingAgent, OutreachAgent ) # Get patent path from task description or metadata # For demo, we'll use a mock patent patent_path = state.get("input_data", {}).get("patent_path", "mock_patent.txt") try: # STEP 1: Document Analysis logger.info("📄 Step 1/4: Analyzing patent document...") doc_agent = DocumentAnalysisAgent( llm_client=self.llm_client, memory_agent=self.memory_agent, vision_ocr_agent=self.vision_ocr_agent ) patent_analysis = await doc_agent.analyze_patent(patent_path) state["agent_outputs"]["document_analysis"] = patent_analysis.model_dump() logger.success(f"✅ Patent analyzed: {patent_analysis.title}") # STEP 2: Market Analysis logger.info("📊 Step 2/4: Analyzing market opportunities...") market_agent = MarketAnalysisAgent( llm_client=self.llm_client, memory_agent=self.memory_agent ) market_analysis = await market_agent.analyze_market(patent_analysis) state["agent_outputs"]["market_analysis"] = market_analysis.model_dump() logger.success(f"✅ Market analyzed: {len(market_analysis.opportunities)} opportunities") # STEP 3: Stakeholder Matching logger.info("🤝 Step 3/4: Finding potential partners...") matching_agent = MatchmakingAgent( llm_client=self.llm_client, memory_agent=self.memory_agent ) matches = await matching_agent.find_matches( patent_analysis, market_analysis, max_matches=10 ) state["agent_outputs"]["matches"] = [m.model_dump() for m in matches] logger.success(f"✅ Found {len(matches)} potential partners") # STEP 4: Generate Valorization Brief logger.info("📝 Step 4/4: Creating valorization brief...") outreach_agent = OutreachAgent( llm_client=self.llm_client, memory_agent=self.memory_agent ) brief = await outreach_agent.create_valorization_brief( patent_analysis, market_analysis, matches ) state["agent_outputs"]["brief"] = brief.model_dump() state["final_output"] = brief.content logger.success(f"✅ Brief created: {brief.pdf_path}") # Set overall execution result state["agent_outputs"]["executor"] = { "result": f"Patent Wake-Up workflow completed successfully", "patent_title": patent_analysis.title, "opportunities_found": len(market_analysis.opportunities), "matches_found": len(matches), "brief_path": brief.pdf_path, "agents_used": ["DocumentAnalysisAgent", "MarketAnalysisAgent", "MatchmakingAgent", "OutreachAgent"], } logger.success("✅ Patent Wake-Up pipeline completed successfully!") except Exception as e: logger.error(f"Patent Wake-Up pipeline failed: {e}") state["agent_outputs"]["executor"] = { "result": f"Pipeline failed: {str(e)}", "error": str(e), "agents_used": [], } state["final_output"] = f"Error: {str(e)}" return state async def _critic_node(self, state: AgentState) -> AgentState: logger.info(f"CRITIC node validating output") state["status"] = TaskStatus.VALIDATING state["current_agent"] = "CriticAgent" if self.critic_agent: from ..agents.base_agent import Task task = Task( id=state["task_id"], description=state["task_description"], metadata={ "output_to_validate": state["final_output"], "output_type": self._get_output_type(state["scenario"]) } ) result_task = await self.critic_agent.process_task(task) if result_task.status == "completed": validation = result_task.result state["validation_score"] = validation.overall_score state["validation_feedback"] = self.critic_agent.get_feedback_for_iteration(validation) state["validation_issues"] = validation.issues state["validation_suggestions"] = validation.suggestions feedback_msg = AIMessage( content=f"Validation score: {validation.overall_score:.2f}\n{state['validation_feedback']}" ) state["messages"].append(feedback_msg) else: llm = self.llm_client.get_llm(complexity="analysis") validation_prompt = HumanMessage( content=f"Validate the following output:\n\n{state['final_output']}\n\nProvide a quality score (0.0-1.0) and feedback." ) response = await llm.ainvoke([validation_prompt]) state["messages"].append(response) state["validation_score"] = 0.90 state["validation_feedback"] = response.content state["validation_issues"] = [] state["validation_suggestions"] = [] logger.info(f"Validation completed: score={state['validation_score']:.2f}") return state async def _refine_node(self, state: AgentState) -> AgentState: logger.info(f"REFINE node preparing for iteration {state['iteration_count'] + 1}") state["status"] = TaskStatus.REFINING state["current_agent"] = "Refiner" state["iteration_count"] += 1 refine_msg = HumanMessage( content=f"Iteration {state['iteration_count']}: Address the following issues:\n{state['validation_feedback']}" ) state["messages"].append(refine_msg) state["intermediate_results"].append({ "iteration": state["iteration_count"] - 1, "output": state["final_output"], "score": state["validation_score"], "feedback": state["validation_feedback"], }) logger.info(f"Refinement prepared for iteration {state['iteration_count']}") return state async def _finish_node(self, state: AgentState) -> AgentState: logger.info(f"FINISH node completing workflow") state["status"] = TaskStatus.COMPLETED state["current_agent"] = None state["success"] = True state["end_time"] = datetime.now() state["execution_time_seconds"] = (state["end_time"] - state["start_time"]).total_seconds() # Store episode in memory for future learning if self.memory_agent and state.get("validation_score", 0) >= 0.75: try: logger.info("Storing episode in memory...") await self.memory_agent.store_episode( task_id=state["task_id"], task_description=state["task_description"], scenario=state["scenario"], workflow_steps=state.get("subtasks", []), outcome={ "final_output": state["final_output"], "validation_score": state.get("validation_score", 0), "success": state["success"], "tools_used": state.get("agent_outputs", {}).get("executor", {}).get("tools_called", []), }, quality_score=state.get("validation_score", 0), execution_time=state["execution_time_seconds"], iterations_used=state.get("iteration_count", 0), ) logger.info(f"Episode stored: {state['task_id']}") except Exception as e: logger.warning(f"Failed to store episode: {e}") completion_msg = AIMessage( content=f"Workflow completed successfully in {state['execution_time_seconds']:.2f}s" ) state["messages"].append(completion_msg) logger.info(f"Workflow completed: {state['task_id']}") return state def _should_refine(self, state: AgentState) -> Literal["refine", "finish"]: score = state.get("validation_score", 0.0) iterations = state.get("iteration_count", 0) if score >= self.quality_threshold: logger.info(f"Quality threshold met ({score:.2f} >= {self.quality_threshold}), finishing") return "finish" if iterations >= state.get("max_iterations", self.max_iterations): logger.warning(f"Max iterations reached ({iterations}), finishing anyway") return "finish" logger.info(f"Refining (score={score:.2f}, iteration={iterations})") return "refine" def _get_scenario_agents(self, scenario: ScenarioType) -> list: scenario_map = { ScenarioType.PATENT_WAKEUP: ["DocumentAnalysisAgent", "MarketAnalysisAgent", "MatchmakingAgent", "OutreachAgent"], ScenarioType.AGREEMENT_SAFETY: ["LegalAnalysisAgent", "ComplianceAgent", "RiskAssessmentAgent", "RecommendationAgent"], ScenarioType.PARTNER_MATCHING: ["ProfilingAgent", "SemanticMatchingAgent", "NetworkAnalysisAgent", "ConnectionFacilitatorAgent"], ScenarioType.GENERAL: ["ExecutorAgent"] } return scenario_map.get(scenario, ["ExecutorAgent"]) def _get_output_type(self, scenario: ScenarioType) -> str: type_map = { ScenarioType.PATENT_WAKEUP: "patent_analysis", ScenarioType.AGREEMENT_SAFETY: "legal_review", ScenarioType.PARTNER_MATCHING: "stakeholder_matching", ScenarioType.GENERAL: "general" } return type_map.get(scenario, "general") async def run( self, task_description: str, scenario: ScenarioType = ScenarioType.GENERAL, task_id: Optional[str] = None, input_data: Optional[Dict[str, Any]] = None, config: Optional[Dict[str, Any]] = None, ) -> WorkflowOutput: if task_id is None: task_id = f"task_{hash(task_description) % 100000}" initial_state = create_initial_state( task_id=task_id, task_description=task_description, scenario=scenario, max_iterations=self.max_iterations, input_data=input_data, ) logger.info(f"Starting workflow for task: {task_id}") try: final_state = await self.app.ainvoke( initial_state, config=config or {"configurable": {"thread_id": task_id}} ) output = state_to_output(final_state) logger.info(f"Workflow completed successfully: {task_id}") return output except Exception as e: logger.error(f"Workflow failed: {e}") initial_state["status"] = TaskStatus.FAILED initial_state["success"] = False initial_state["error"] = str(e) initial_state["end_time"] = datetime.now() return state_to_output(initial_state) async def stream( self, task_description: str, scenario: ScenarioType = ScenarioType.GENERAL, task_id: Optional[str] = None, config: Optional[Dict[str, Any]] = None, ): if task_id is None: task_id = f"task_{hash(task_description) % 100000}" initial_state = create_initial_state( task_id=task_id, task_description=task_description, scenario=scenario, max_iterations=self.max_iterations, ) async for event in self.app.astream( initial_state, config=config or {"configurable": {"thread_id": task_id}} ): yield event def create_workflow( llm_client: LangChainOllamaClient, planner_agent: Optional[Any] = None, critic_agent: Optional[Any] = None, memory_agent: Optional[Any] = None, vision_ocr_agent: Optional[Any] = None, quality_threshold: float = 0.85, max_iterations: int = 3, ) -> SparknetWorkflow: return SparknetWorkflow( llm_client=llm_client, planner_agent=planner_agent, critic_agent=critic_agent, memory_agent=memory_agent, vision_ocr_agent=vision_ocr_agent, quality_threshold=quality_threshold, max_iterations=max_iterations, )