Spaces:
Sleeping
Sleeping
File size: 1,704 Bytes
d3d1a7c 7a99a44 d3d1a7c 6de5dfa 8297d3e 352d26f 7a99a44 352d26f 16edc8c 352d26f | 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 | from langgraph.graph import START, StateGraph
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from typing import Sequence
from typing_extensions import TypedDict, Annotated
class State(TypedDict):
input: str
chat_history: Annotated[Sequence[BaseMessage], "add_messages"]
context: str
answer: str
# For direct RAG chains
def build_graph(rag_chain):
def call_model(state: State):
response = rag_chain.invoke(state)
return {
"chat_history": [
HumanMessage(state["input"]),
AIMessage(response["answer"]),
],
"context": response.get("context", ""),
"answer": response.get("answer", "")
}
workflow = StateGraph(state_schema=State)
workflow.add_node("model", call_model)
workflow.add_edge(START, "model")
#memory = MemorySaver()
#return workflow.compile(checkpointer=memory)
return workflow.compile() # Stateless; relies on session memory only
# For agent_chain.invoke
def build_graph_with_callable(call_fn):
def call_model(state: State):
response = call_fn({"input": state["input"]})
return {
"chat_history": [
HumanMessage(state["input"]),
AIMessage(response.get("output", response.get("answer", ""))),
],
"context": response.get("context", ""),
"answer": response.get("output", response.get("answer", ""))
}
workflow = StateGraph(state_schema=State)
workflow.add_node("model", call_model)
workflow.add_edge(START, "model")
return workflow.compile()
|