rag-playground/langgraph_client.py at main · kvpratama/rag-playground · GitHub

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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
from raptor.graph import workflow as raptor_workflow
from crag.graph import workflow as crag_workflow
from selfrag.graph import workflow as selfrag_workflow
from agenticrag.graph import workflow as agenticrag_workflow
from adaptiverag.graph import workflow as adaptiverag_workflow
import logging

logger = logging.getLogger(__name__)

def run_raptor(url, max_depth, question, thread_id):
    input_data = {"url": url, "max_depth": max_depth, "question": question}
    result = raptor_workflow.invoke(input_data, config={"configurable": {"thread_id": thread_id}})
    return result["answer"]

def run_crag(urls, question, thread_id):
    input_data = {"urls": urls, "question": question}
    result = crag_workflow.invoke(input_data, config={"configurable": {"thread_id": thread_id}})
    if result:
        return result.get("answer", "")
    return ""

def run_crag_stream(urls, question, thread_id):
    logger.info(f"Running crag with urls: {urls} and question: {question}")
    input_data = {"urls": urls, "question": question}
    for event in crag_workflow.stream(input_data, config={"configurable": {"thread_id": thread_id}}, subgraphs=True, stream_mode="custom"):
        _, data = event  # event[1] → data
        output = data.get("custom_key", "")
        yield output

def run_selfrag(urls, question, thread_id):
    logger.info(f"Running selfrag with urls: {urls} and question: {question}")
    input_data = {"urls": urls, "question": question}
    result = selfrag_workflow.invoke(input_data, config={"configurable": {"thread_id": thread_id}})
    logger.info(f"Result: {result}")
    if result:
        return result.get("generation", "")
    return ""

def run_selfrag_stream(urls, question, thread_id):
    logger.info(f"Running selfrag with urls: {urls} and question: {question}")
    input_data = {"urls": urls, "question": question}
    for event in selfrag_workflow.stream(input_data, config={"configurable": {"thread_id": thread_id}}, subgraphs=True, stream_mode="custom"):
        _, data = event  # event[1] → data
        output = data.get("custom_key", "")
        yield output

def run_agenticrag(urls, question, thread_id):
    logger.info(f"Running agenticrag with urls: {urls} and question: {question}")
    input_data = {"urls": urls, "question": question}
    result = agenticrag_workflow.invoke(input_data, config={"configurable": {"thread_id": thread_id}})
    logger.info(f"Result: {result}")
    if result:
        return result.get("generation", "")
    return ""

def run_agenticrag_stream(urls, question, thread_id):
    logger.info(f"Running agenticrag with urls: {urls} and question: {question}")
    input_data = {"urls": urls, "question": question}
    for event in agenticrag_workflow.stream(input_data, config={"configurable": {"thread_id": thread_id}}, subgraphs=True, stream_mode="custom"):
        _, data = event  # event[1] → data
        output = data.get("custom_key", "")
        yield output

def run_adaptiverag(urls, question, thread_id):
    logger.info(f"Running adaptiverag with urls: {urls} and question: {question}")
    input_data = {"urls": urls, "question": question}
    result = adaptiverag_workflow.invoke(input_data, config={"configurable": {"thread_id": thread_id}})
    logger.info(f"Result: {result}")
    if result:
        return result.get("generation", "")
    return ""

def run_adaptiverag_stream(urls, question, thread_id):
    logger.info(f"Running adaptiverag with urls: {urls} and question: {question}")
    input_data = {"urls": urls, "question": question}
    for event in adaptiverag_workflow.stream(input_data, config={"configurable": {"thread_id": thread_id}}, subgraphs=True, stream_mode="custom"):
        _, data = event  # event[1] → data
        output = data.get("custom_key", "")
        yield output