Spaces:
Sleeping
Sleeping
File size: 1,503 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 57 58 59 60 61 62 | from langgraph.graph import StateGraph, START, END
from src.MultiRag.models.worker_model import State
from src.MultiRag.nodes.worker import (
pdf,
txt,
docs,
image,
url,
decider,
search
)
from src.MultiRag.constants import AVAILABLE_ANALYSIS
import logging
logging.info("Building worker sub graph")
graph = StateGraph(State)
graph.add_node("decider", decider.decider_node)
graph.add_node("pdf", pdf.pdf_node)
graph.add_node("txt", txt.txt_node)
graph.add_node("docs", docs.docs_node)
graph.add_node("url", url.url_node)
graph.add_node("image", image.image_node)
graph.add_node("search", search.search_node)
def route_fn(state: State):
logging.info(f"Routing based on file_type: {state.file_type}")
if state.file_type in AVAILABLE_ANALYSIS:
return state.file_type
return "end"
graph.add_conditional_edges(
START,
route_fn,
{
"pdf": "pdf",
"txt": "txt",
"docs": "docs",
"png": "image",
"url": "url",
"search": "search",
"end":END
}
)
graph.add_edge("pdf", END)
graph.add_edge("txt", END)
graph.add_edge("docs", END)
graph.add_edge("url", END)
graph.add_edge("image", END)
graph.add_edge("search", END)
graph = graph.compile()
try:
with open("worker_sub_graph.png", "wb") as f:
f.write(graph.get_graph().draw_mermaid_png())
logging.info("Graph image saved successfully")
except Exception as e:
logging.error(f"Error saving graph: {e}")
raise Exception(e) |