Spaces:
Sleeping
Sleeping
File size: 1,578 Bytes
2436836 195e631 2436836 195e631 2436836 195e631 2436836 195e631 2436836 195e631 2436836 195e631 2436836 | 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 | # # app/graph/nodes/general_agent.py
# from app.core.llm_engine import llm
# from langchain_core.output_parsers import StrOutputParser
# from app.core.prompts.general_prompt import general_prompt
# def general_agent_node(state):
# query = state.get("query")
# chain = general_prompt | llm | StrOutputParser()
# response = chain.invoke({"query": query})
# return {
# **state,
# "general_answer": response.strip()
# }
from langchain_core.output_parsers import StrOutputParser
from app.core.prompts.general_prompt import general_prompt
from app.core.llm_engine import llm, get_streaming_llm
# -------------------------------------------------------
# Existing synchronous node
# -------------------------------------------------------
def general_agent_node(state):
query = state.get("query", "")
chain = general_prompt | llm | StrOutputParser()
response = chain.invoke({
"query": query
})
return {
**state,
"general_answer": response.strip()
}
# -------------------------------------------------------
# NEW
# Streaming version
# -------------------------------------------------------
async def general_agent_stream(state):
query = state.get("query", "")
stream_llm = get_streaming_llm()
chain = general_prompt | stream_llm
async for chunk in chain.astream({
"query": query
}):
if hasattr(chunk, "content") and chunk.content:
yield chunk.content
|