Files
langgraph/examples/runnable-pubsub.py
T
2023-08-18 15:10:30 +01:00

139 lines
4.1 KiB
Python

from operator import itemgetter
from pprint import pprint
from langchain.chat_models.openai import ChatOpenAI
from langchain.prompts import SystemMessagePromptTemplate
from langchain.runnables.openai_functions import OpenAIFunctionsRouter
from langchain.schema.output_parser import StrOutputParser
from permchain.connection_inmemory import InMemoryPubSubConnection
from permchain.pubsub import PubSub
from permchain.topic import Topic
drafter_prompt = (
SystemMessagePromptTemplate.from_template(
"You are an expert on turtles, who likes to write in pirate-speak. You have been tasked by your editor with drafting a 100-word article answering the following question."
)
+ "Question:\n\n{question}"
)
reviser_prompt = (
SystemMessagePromptTemplate.from_template(
"You are an expert on turtles. You have been tasked by your editor with revising the following draft, which was written by a non-expert. You may follow the editor's notes or not, as you see fit."
)
+ "Draft:\n\n{draft}"
+ "Editor's notes:\n\n{notes}"
)
editor_prompt = (
SystemMessagePromptTemplate.from_template(
"You are an editor. You have been tasked with editing the following draft, which was written by a non-expert. Please accept the draft if it is good enough to publish, or send it for revision, along with your notes to guide the revision."
)
+ "Draft:\n\n{draft}"
)
drafter_llm = ChatOpenAI(model="gpt-3.5-turbo")
editor_llm = ChatOpenAI(model="gpt-4")
# create topics
editor_inbox = Topic("editor_inbox")
reviser_inbox = Topic("reviser_inbox")
# write a first draft
drafter = (
Topic.IN.subscribe()
| {"draft": drafter_prompt | drafter_llm | StrOutputParser()}
| editor_inbox.publish()
)
# edit every draft, produce revision notes or accept
editor = (
editor_inbox.subscribe()
| editor_prompt
| editor_llm.bind(
functions=[
{
"name": "revise",
"description": "Sends the draft for revision",
"parameters": {
"type": "object",
"properties": {
"notes": {
"type": "string",
"description": "The editor's notes to guide the revision.",
},
},
},
},
{
"name": "accept",
"description": "Accepts the draft",
"parameters": {
"type": "object",
"properties": {"ready": {"const": True}},
},
},
]
)
| OpenAIFunctionsRouter(
{
"revise": (
{
"notes": itemgetter("notes"),
"draft": editor_inbox.current() | itemgetter("draft"),
"question": Topic.IN.current() | itemgetter("question"),
}
| reviser_inbox.publish()
),
"accept": editor_inbox.current() | Topic.OUT.publish(),
},
)
)
# every time revision notes are posted, revise latest draft
reviser = (
reviser_inbox.subscribe()
| {"draft": reviser_prompt | drafter_llm | StrOutputParser()}
| editor_inbox.publish()
)
web_researcher = PubSub(
processes=(drafter, editor, reviser),
connection=InMemoryPubSubConnection(),
)
# for output in web_researcher.stream({"question": "What food do turtles eat?"}):
# print("got output", output)
# print("---done with stream()---")
# pprint(web_researcher.invoke({"question": "What food do turtles eat?"}))
pprint(
web_researcher.batch(
[{"question": "What food do turtles eat?"}, {"question": "What is art?"}]
)
)
# agent = PubSub(
# Channel.IN | Channel("planner"),
# Channel("executor") | executor | Channel("planner"),
# Channel("planner")
# | planner
# | {"action": Channel("executor"), "finish": Channel.OUT},
# )
# graph = (
# drafter
# | editor
# | RouterRunnable(
# {
# "send_for_revision": reviser,
# "accept_draft": lambda x: x["draft"],
# }
# )
# )