Classification langchain example

examples/langchain-langserve/app/classification_chain.py

@@ -0,0 +1,132 @@
+import os
+
+from langchain.callbacks.tracers import ConsoleCallbackHandler
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts.prompt import PromptTemplate
+from langchain_core.pydantic_v1 import BaseModel
+from langchain_core.runnables import (
+    RunnableLambda,
+    RunnablePassthrough,
+    RunnableBranch
+)
+from langchain_core.runnables.history import RunnableWithMessageHistory
+
+from zep_python import ZepClient
+from zep_python.langchain import ZepChatMessageHistory
+
+from langchain_community.chat_models import ChatAnthropic
+
+ZEP_API_KEY = os.environ.get("ZEP_API_KEY")  # Required for Zep Cloud
+ZEP_API_URL = os.environ.get(
+    "ZEP_API_URL"
+)  # only required if you're using Zep Open Source
+
+if ZEP_API_KEY is None:
+    raise ValueError(
+        "ZEP_API_KEY is required for Zep Cloud. "
+        "Remove this check if using Zep Open Source."
+    )
+
+zep = ZepClient(
+    api_key=ZEP_API_KEY,
+    api_url=ZEP_API_URL,  # only required if you're using Zep Open Source
+)
+
+langchain_chain = (
+        PromptTemplate.from_template(
+            """You are an expert in langchain. \
+    Always answer questions starting with "As Harrison Chase told me". \
+    Respond to the following question:
+
+    Question: {question}
+    Answer:"""
+        )
+        | ChatAnthropic()
+)
+anthropic_chain = (
+        PromptTemplate.from_template(
+            """You are an expert in anthropic. \
+    Always answer questions starting with "As Dario Amodei told me". \
+    Respond to the following question:
+
+    Question: {question}
+    Answer:"""
+        )
+        | ChatAnthropic()
+)
+general_chain = (
+        PromptTemplate.from_template(
+            """Respond to the following question:
+
+    Question: {question}
+    Answer:"""
+        )
+        | ChatAnthropic()
+)
+
+branch = RunnableBranch(
+    (lambda x: "anthropic" in x["topic"].lower(), anthropic_chain),
+    (lambda x: "langchain" in x["topic"].lower(), langchain_chain),
+    general_chain,
+)
+
+topic_classifier = (
+        PromptTemplate.from_template(
+            """Given the user question below, classify it as either being about `LangChain`, `Anthropic`, or `Other`.
+
+    Do not respond with more than one word.
+
+    <question>
+    {question}
+    </question>
+
+    Classification:"""
+        )
+        | ChatAnthropic()
+        | StrOutputParser()
+)
+
+
+# User input
+class UserInput(BaseModel):
+    question: str
+    session_id: str
+
+
+def classify_session(session_id: str):
+    result = zep.memory.classify_session(
+        session_id,
+        "tool",
+        ["langchain", "anthropic", "none"],
+        last_n=1,
+        persist=False
+    )
+    return result.class_
+
+
+def invoke_chain(user_input: UserInput):
+    result_chain = RunnableWithMessageHistory(
+        RunnablePassthrough.assign(session_id=lambda x: user_input["session_id"])
+        | {
+            "question": lambda x: x["question"],
+            "topic": lambda x: classify_session(x["session_id"]),
+        } | branch,
+        lambda session_id: ZepChatMessageHistory(
+            session_id=session_id,
+            zep_client=zep,
+            memory_type="perpetual",
+        ),
+        input_messages_key="question",
+        history_messages_key="chat_history",
+    )
+
+    return result_chain.invoke(
+        user_input, config={"configurable": {"session_id": user_input["session_id"]}}
+    )
+
+
+chain = (
+    RunnableLambda(invoke_chain)
+    .with_types(input_type=UserInput)
+    .with_config(callbacks=[ConsoleCallbackHandler()])
+)

examples/langchain-langserve/app/server.py

@@ -10,6 +10,7 @@
     chain as message_history_vector_store_chain,
 )
 from rag_vector_store_chain import chain as rag_vector_store_chain
+from classification_chain import chain as classification_chain
 
 app = FastAPI()
 
@@ -24,6 +25,7 @@ async def redirect_root_to_docs():
 add_routes(
     app, message_history_vector_store_chain, path="/message_history_vector_store"
 )
+add_routes(app, classification_chain, path="/classification")
 
 if __name__ == "__main__":
     import uvicorn