Spaces:
Sleeping
Sleeping
File size: 2,542 Bytes
5551822 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | import logging
from utils.asyncHandler import asyncHandler
from src.MultiRag.models.rag_model import State
from src.MultiRag.models.orchestrator_output_model import OrchestratorOutput
from src.MultiRag.llm.llm_loader import llm
from src.MultiRag.prompts.prompt_templates import ORCHESTRATOR_PROMPT
from langchain_core.messages import SystemMessage, HumanMessage
@asyncHandler
async def orchestrator_node(state:State):
logging.info("Entered in the orchestrator_node")
logging.info(f"Current messages: {len(state.get('messages', []))} message(s)")
orchestrator_llm = llm.with_structured_output(OrchestratorOutput, method="tool_call")
user_content = state.get('userContent', [])
files_info = "\n".join([f"- Name: {c.name}, Path: {c.path}, About: {c.about}" for c in user_content])
logging.info(f"Files available for orchestration: {len(user_content)}")
system_prompt = ORCHESTRATOR_PROMPT + f"\n\n### Available Files:\n{files_info}\n\nWhen using a worker, you MUST specify the exact 'file_path' and 'file_type' (one of: pdf, txt, docs, png, url) from the list above."
prompt= [SystemMessage(content=system_prompt)]+ state.get('messages', [])
logging.info("Invoking orchestrator LLM with file context...")
try:
response = await orchestrator_llm.ainvoke(prompt)
except Exception as e:
logging.error(f"Error in orchestrator ainvoke: {e}")
response = None
if response is None:
logging.info("Structured output failed, attempting manual JSON parsing...")
raw_res = await llm.ainvoke(prompt)
import json
import re
content = raw_res.content
logging.info(f"Raw orchestrator response: {content}")
json_match = re.search(r'\{.*\}', content, re.DOTALL)
if json_match:
try:
json_data = json.loads(json_match.group())
response = OrchestratorOutput(**json_data)
logging.info("Successfully parsed JSON manually.")
except Exception as e:
logging.error(f"Manual JSON parsing failed: {e}")
else:
logging.warning("No JSON block found in orchestrator response. Attempting to construct plan from text...")
if "worker" in content.lower() and "docs/" in content:
logging.info("Detected worker mention in text, but couldn't parse JSON. This model might need a better prompt.")
logging.info(f"Final plan decided: {response}")
return {"plan": response}
|