mirror of
https://github.com/hwchase17/langchain
synced 2024-11-08 07:10:35 +00:00
114 lines
3.2 KiB
Python
114 lines
3.2 KiB
Python
from operator import itemgetter
|
|
from typing import Literal
|
|
|
|
from langchain.chat_models import ChatOpenAI
|
|
from langchain.output_parsers.openai_functions import PydanticAttrOutputFunctionsParser
|
|
from langchain.prompts import ChatPromptTemplate
|
|
from langchain.pydantic_v1 import BaseModel, Field
|
|
from langchain.retrievers import (
|
|
ArxivRetriever,
|
|
KayAiRetriever,
|
|
PubMedRetriever,
|
|
WikipediaRetriever,
|
|
)
|
|
from langchain.schema import StrOutputParser
|
|
from langchain.schema.runnable import (
|
|
RouterRunnable,
|
|
RunnableParallel,
|
|
RunnablePassthrough,
|
|
)
|
|
from langchain.utils.openai_functions import convert_pydantic_to_openai_function
|
|
|
|
pubmed = PubMedRetriever(top_k_results=5).with_config(run_name="pubmed")
|
|
arxiv = ArxivRetriever(top_k_results=5).with_config(run_name="arxiv")
|
|
sec = KayAiRetriever.create(
|
|
dataset_id="company", data_types=["10-K"], num_contexts=5
|
|
).with_config(run_name="sec_filings")
|
|
wiki = WikipediaRetriever(top_k_results=5, doc_content_chars_max=2000).with_config(
|
|
run_name="wiki"
|
|
)
|
|
|
|
llm = ChatOpenAI(model="gpt-3.5-turbo")
|
|
|
|
|
|
class Search(BaseModel):
|
|
"""Search for relevant documents by question topic."""
|
|
|
|
question_resource: Literal[
|
|
"medical paper", "scientific paper", "public company finances report", "general"
|
|
] = Field(
|
|
...,
|
|
description=(
|
|
"The type of resource that would best help answer the user's question. "
|
|
"If none of the types are relevant return 'general'."
|
|
),
|
|
)
|
|
|
|
|
|
retriever_name = {
|
|
"medical paper": "PubMed",
|
|
"scientific paper": "ArXiv",
|
|
"public company finances report": "SEC filings (Kay AI)",
|
|
"general": "Wikipedia",
|
|
}
|
|
|
|
classifier = (
|
|
llm.bind(
|
|
functions=[convert_pydantic_to_openai_function(Search)],
|
|
function_call={"name": "Search"},
|
|
)
|
|
| PydanticAttrOutputFunctionsParser(
|
|
pydantic_schema=Search, attr_name="question_resource"
|
|
)
|
|
| retriever_name.get
|
|
)
|
|
|
|
retriever_map = {
|
|
"PubMed": pubmed,
|
|
"ArXiv": arxiv,
|
|
"SEC filings (Kay AI)": sec,
|
|
"Wikipedia": wiki,
|
|
}
|
|
router_retriever = RouterRunnable(runnables=retriever_map)
|
|
|
|
|
|
def format_docs(docs):
|
|
return "\n\n".join(f"Source {i}:\n{doc.page_content}" for i, doc in enumerate(docs))
|
|
|
|
|
|
system = """Answer the user question. Use the following sources to help \
|
|
answer the question. If you don't know the answer say "I'm not sure, I couldn't \
|
|
find information on {{topic}}."
|
|
|
|
Sources:
|
|
|
|
{sources}"""
|
|
prompt = ChatPromptTemplate.from_messages([("system", system), ("human", "{question}")])
|
|
|
|
|
|
class Question(BaseModel):
|
|
__root__: str
|
|
|
|
|
|
retriever_chain = (
|
|
{"input": itemgetter("question"), "key": itemgetter("retriever_choice")}
|
|
| router_retriever
|
|
| format_docs
|
|
).with_config(run_name="retrieve")
|
|
answer_chain = (
|
|
{"sources": retriever_chain, "question": itemgetter("question")}
|
|
| prompt
|
|
| llm
|
|
| StrOutputParser()
|
|
)
|
|
chain = (
|
|
(
|
|
RunnableParallel(
|
|
question=RunnablePassthrough(), retriever_choice=classifier
|
|
).with_config(run_name="classify")
|
|
| RunnablePassthrough.assign(answer=answer_chain).with_config(run_name="answer")
|
|
)
|
|
.with_config(run_name="QA with router")
|
|
.with_types(input_type=Question)
|
|
)
|