open-notebook/open_notebook/graphs/prompt.py
LUIS NOVO ab5560c9a2 refactor: reorganize folder structure for better maintainability
Changes:
- Move migrations/ under open_notebook/database/migrations/
- Extract AI models to open_notebook/ai/ (Model, ModelManager, provision)
- Extract podcasts to open_notebook/podcasts/ (EpisodeProfile, SpeakerProfile, PodcastEpisode)
- Reorganize prompts to mirror graphs structure (chat/, source_chat/)

This improves code organization by:
- Consolidating database concerns (migrations now with database code)
- Separating AI infrastructure from domain entities
- Isolating podcast feature into its own module
- Creating consistent prompt/graph naming conventions

All 52 tests pass.
2026-01-03 14:04:27 -03:00

42 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 typing_extensions import TypedDict
from open_notebook.ai.provision 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) # type: ignore[type-var]
agent_state.add_edge(START, "agent")
agent_state.add_edge("agent", END)
graph = agent_state.compile()