open-notebook/open_notebook/graphs/prompt.py
Luis Novo d7b0fff954
Api podcast migration (#93)
Creates the API layer for Open Notebook
Creates a services API gateway for the Streamlit front-end
Migrates the SurrealDB SDK to the official one
Change all database calls to async
New podcast framework supporting multiple speaker configurations
Implement the surreal-commands library for async processing
Improve docker image and docker-compose configurations
2025-07-17 08:36:11 -03:00

43 lines
1.2 KiB
Python

from typing import Any, Optional
from ai_prompter import Prompter
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, START, StateGraph
from loguru import logger
from typing_extensions import TypedDict
from open_notebook.graphs.utils import provision_langchain_model
class PatternChainState(TypedDict):
prompt: str
parser: Optional[Any]
input_text: str
output: str
async def call_model(state: dict, config: RunnableConfig) -> dict:
content = state["input_text"]
system_prompt = Prompter(
template_text=state["prompt"], parser=state.get("parser")
).render(data=state)
payload = [SystemMessage(content=system_prompt)] + [HumanMessage(content=content)]
chain = await provision_langchain_model(
str(payload),
config.get("configurable", {}).get("model_id"),
"transformation",
max_tokens=5000,
)
response = await chain.ainvoke(payload)
return {"output": response.content}
agent_state = StateGraph(PatternChainState)
agent_state.add_node("agent", call_model)
agent_state.add_edge(START, "agent")
agent_state.add_edge("agent", END)
graph = agent_state.compile()