langchain

core[patch]: fix no current event loop for sql history in async mode

#22933

Merged

from typing import AsyncIterator, Any

from fastapi import APIRouter, FastAPI
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories.sql import SQLChatMessageHistory
from langchain_core.runnables import Runnable
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langchain_core.tools import Tool
from langchain_openai import AzureChatOpenAI
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.prompts import (
    ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate,
    SystemMessagePromptTemplate
)
from langchain.schema.runnable import RunnableLambda
from langserve import add_routes
from langserve.pydantic_v1 import BaseModel, Field
from sqlalchemy.ext.asyncio import create_async_engine


# Config
sql_server_url = "mysql+asyncmy://dbuser:dbpass@localhost:3306/history"

azure_endpoint = "fixme"
azure_deployment = "fixme"
openai_api_version = "fixme"
openai_api_key = "fixme"


# Model
class LLMInput(BaseModel):
    question: str = Field(
        ...,
        description="user question",
    )
    system_prompt: str = Field(
        default="You are a helpful AI assistant",
        description="system prompt",
    )


# History Stuff
async_sql_engine = create_async_engine(sql_server_url, pool_recycle=600)

history_factory_config = [
    ConfigurableFieldSpec(
        id="user_id",
        annotation=str,
        name="User ID",
        description="Unique identifier for the user.",
        default="",
        is_shared=True,
    ),
    ConfigurableFieldSpec(
        id="conversation_id",
        annotation=str,
        name="Conversation ID",
        description="Unique identifier for the conversation.",
        default="",
        is_shared=True,
    ),
]

def get_session_history(
    user_id: str,
    conversation_id: str,
) -> BaseChatMessageHistory:
    """Get a chat history from a session ID."""
    session_id = f"{user_id}_{conversation_id}"

    return SQLChatMessageHistory(
        session_id=session_id,
        connection=async_sql_engine,
    )


# Tool Stuff
def search(query: str, **kwargs) -> str:
    return f"{query} found"


class SearchToolInput(BaseModel):
    """Input for music tool."""
    query: str = Field(description="keyword for search")


search_tool = Tool(
    name="web_search",
    description="web_search is a tool to search online information.",
    func=search,
    args_schema=SearchToolInput,
)


def build_llm_chain(llm_input: LLMInput) -> Runnable:
    llm = AzureChatOpenAI(
        azure_endpoint=azure_endpoint,
        azure_deployment=azure_deployment,
        openai_api_version=openai_api_version,
        openai_api_key=openai_api_key,
        streaming=True,
        temperature=0.01,
        max_tokens=4000,
    )
    tools = [search_tool]
    prompt = ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template(llm_input.system_prompt),
        MessagesPlaceholder(variable_name="history"),
        HumanMessagePromptTemplate.from_template("{question}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ])
    agent = create_tool_calling_agent(
        llm, tools, prompt,
    )
    return AgentExecutor(
        agent=agent,
        tools=tools,
    )


def build_chain(llm_input: LLMInput) -> Runnable:
    llm_chain = build_llm_chain(llm_input)

    return RunnableWithMessageHistory(
        llm_chain,
        get_session_history,
        input_messages_key="question",
        history_messages_key="history",
        history_factory_config=history_factory_config
    ).with_types(
        input_type=LLMInput
    )

async def llm_main(_input: dict[str, Any]) -> AsyncIterator[Any]:
    _input["question"] = _input["question"].lstrip("\n").rstrip()
    _input["system_prompt"] = _input["system_prompt"] or "You are a helpful AI assistant"

    llm_input: LLMInput = LLMInput.parse_obj(_input)

    chain = build_chain(llm_input)
    config = {
        "configurable": {
            "user_id": "user1",
            "conversation_id": "conv1",
        }
    }

    async for event in chain.astream_events(_input, config, version="v1"):
        yield event


router = APIRouter()

add_routes(
    router,
    RunnableLambda(llm_main).with_types(input_type=LLMInput),
    enabled_endpoints=["invoke", "batch", "stream"]
)

app = FastAPI(
    title="Test Server",
    version="0.01",
    description="Test Server Powered by LangChain&LLM",
)
app.include_router(router)

from langserve import RemoteRunnable

client = RemoteRunnable('http://localhost:8000/')
for chunk in yuan.stream({'question': "What's langchain"}):
    print(chunk)