mirror of
https://github.com/hwchase17/langchain
synced 2024-11-08 07:10:35 +00:00
102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
import os
|
|
from operator import itemgetter
|
|
from typing import List, Tuple
|
|
|
|
from langchain.chat_models import ChatOpenAI
|
|
from langchain.embeddings import OpenAIEmbeddings
|
|
from langchain.retrievers import SelfQueryRetriever
|
|
from langchain.schema import format_document
|
|
from langchain.schema.output_parser import StrOutputParser
|
|
from langchain.schema.runnable import RunnableParallel, RunnablePassthrough
|
|
from langchain.vectorstores.elasticsearch import ElasticsearchStore
|
|
from pydantic.v1 import BaseModel, Field
|
|
|
|
from .prompts import CONDENSE_QUESTION_PROMPT, DOCUMENT_PROMPT, LLM_CONTEXT_PROMPT
|
|
|
|
ELASTIC_CLOUD_ID = os.getenv("ELASTIC_CLOUD_ID")
|
|
ELASTIC_USERNAME = os.getenv("ELASTIC_USERNAME", "elastic")
|
|
ELASTIC_PASSWORD = os.getenv("ELASTIC_PASSWORD")
|
|
ES_URL = os.getenv("ES_URL", "http://localhost:9200")
|
|
ELASTIC_INDEX_NAME = os.getenv("ELASTIC_INDEX_NAME", "workspace-search-example")
|
|
|
|
if ELASTIC_CLOUD_ID and ELASTIC_USERNAME and ELASTIC_PASSWORD:
|
|
es_connection_details = {
|
|
"es_cloud_id": ELASTIC_CLOUD_ID,
|
|
"es_user": ELASTIC_USERNAME,
|
|
"es_password": ELASTIC_PASSWORD,
|
|
}
|
|
else:
|
|
es_connection_details = {"es_url": ES_URL}
|
|
|
|
vecstore = ElasticsearchStore(
|
|
ELASTIC_INDEX_NAME,
|
|
embedding=OpenAIEmbeddings(),
|
|
**es_connection_details,
|
|
)
|
|
|
|
document_contents = "The purpose and specifications of a workplace policy."
|
|
metadata_field_info = [
|
|
{"name": "name", "type": "string", "description": "Name of the workplace policy."},
|
|
{
|
|
"name": "created_on",
|
|
"type": "date",
|
|
"description": "The date the policy was created in ISO 8601 date format (YYYY-MM-DD).", # noqa: E501
|
|
},
|
|
{
|
|
"name": "updated_at",
|
|
"type": "date",
|
|
"description": "The date the policy was last updated in ISO 8601 date format (YYYY-MM-DD).", # noqa: E501
|
|
},
|
|
{
|
|
"name": "location",
|
|
"type": "string",
|
|
"description": "Where the policy text is stored. The only valid values are ['github', 'sharepoint'].", # noqa: E501
|
|
},
|
|
]
|
|
llm = ChatOpenAI(temperature=0)
|
|
retriever = SelfQueryRetriever.from_llm(
|
|
llm, vecstore, document_contents, metadata_field_info
|
|
)
|
|
|
|
|
|
def _combine_documents(docs: List) -> str:
|
|
return "\n\n".join(format_document(doc, prompt=DOCUMENT_PROMPT) for doc in docs)
|
|
|
|
|
|
def _format_chat_history(chat_history: List[Tuple]) -> str:
|
|
return "\n".join(f"Human: {human}\nAssistant: {ai}" for human, ai in chat_history)
|
|
|
|
|
|
class InputType(BaseModel):
|
|
question: str
|
|
chat_history: List[Tuple[str, str]] = Field(default_factory=list)
|
|
|
|
|
|
standalone_question = (
|
|
{
|
|
"question": itemgetter("question"),
|
|
"chat_history": lambda x: _format_chat_history(x["chat_history"]),
|
|
}
|
|
| CONDENSE_QUESTION_PROMPT
|
|
| llm
|
|
| StrOutputParser()
|
|
)
|
|
|
|
|
|
def route_question(input):
|
|
if input.get("chat_history"):
|
|
return standalone_question
|
|
else:
|
|
return RunnablePassthrough()
|
|
|
|
|
|
_context = RunnableParallel(
|
|
context=retriever | _combine_documents,
|
|
question=RunnablePassthrough(),
|
|
)
|
|
|
|
|
|
chain = (
|
|
standalone_question | _context | LLM_CONTEXT_PROMPT | llm | StrOutputParser()
|
|
).with_types(input_type=InputType)
|