mirror of
https://github.com/hwchase17/langchain
synced 2024-11-18 09:25:54 +00:00
da94c750c5
<!-- Thank you for contributing to LangChain! Replace this entire comment with: - **Description:** a description of the change, - **Issue:** the issue # it fixes (if applicable), - **Dependencies:** any dependencies required for this change, - **Tag maintainer:** for a quicker response, tag the relevant maintainer (see below), - **Twitter handle:** we announce bigger features on Twitter. If your PR gets announced, and you'd like a mention, we'll gladly shout you out! Please make sure your PR is passing linting and testing before submitting. Run `make format`, `make lint` and `make test` to check this locally. See contribution guidelines for more information on how to write/run tests, lint, etc: https://github.com/langchain-ai/langchain/blob/master/.github/CONTRIBUTING.md If you're adding a new integration, please include: 1. a test for the integration, preferably unit tests that do not rely on network access, 2. an example notebook showing its use. It lives in `docs/extras` directory. If no one reviews your PR within a few days, please @-mention one of @baskaryan, @eyurtsev, @hwchase17. --> --------- Co-authored-by: Matvey Arye <mat@timescale.com>
113 lines
3.2 KiB
Python
113 lines
3.2 KiB
Python
# ruff: noqa: E501
|
|
|
|
import os
|
|
from datetime import timedelta
|
|
|
|
from langchain.chains.query_constructor.base import AttributeInfo
|
|
from langchain.chat_models import ChatOpenAI
|
|
from langchain.embeddings.openai import OpenAIEmbeddings
|
|
from langchain.llms import OpenAI
|
|
from langchain.prompts import ChatPromptTemplate
|
|
from langchain.retrievers.self_query.base import SelfQueryRetriever
|
|
from langchain.schema.output_parser import StrOutputParser
|
|
from langchain.schema.runnable import RunnableParallel, RunnablePassthrough
|
|
from langchain.vectorstores.timescalevector import TimescaleVector
|
|
from pydantic import BaseModel
|
|
|
|
from .load_sample_dataset import load_ts_git_dataset
|
|
|
|
# to enable debug uncomment the following lines:
|
|
# from langchain.globals import set_debug
|
|
# set_debug(True)
|
|
|
|
# from dotenv import find_dotenv, load_dotenv
|
|
# _ = load_dotenv(find_dotenv())
|
|
|
|
if os.environ.get("TIMESCALE_SERVICE_URL", None) is None:
|
|
raise Exception("Missing `TIMESCALE_SERVICE_URL` environment variable.")
|
|
|
|
SERVICE_URL = os.environ["TIMESCALE_SERVICE_URL"]
|
|
LOAD_SAMPLE_DATA = os.environ.get("LOAD_SAMPLE_DATA", False)
|
|
|
|
|
|
# DATASET SPECIFIC CODE
|
|
# Load the sample dataset. You will have to change this to load your own dataset.
|
|
collection_name = "timescale_commits"
|
|
partition_interval = timedelta(days=7)
|
|
if LOAD_SAMPLE_DATA:
|
|
load_ts_git_dataset(
|
|
SERVICE_URL,
|
|
collection_name=collection_name,
|
|
num_records=500,
|
|
partition_interval=partition_interval,
|
|
)
|
|
|
|
# This will change depending on the metadata stored in your dataset.
|
|
document_content_description = "The git log commit summary containing the commit hash, author, date of commit, change summary and change details"
|
|
metadata_field_info = [
|
|
AttributeInfo(
|
|
name="id",
|
|
description="A UUID v1 generated from the date of the commit",
|
|
type="uuid",
|
|
),
|
|
AttributeInfo(
|
|
# This is a special attribute represent the timestamp of the uuid.
|
|
name="__uuid_timestamp",
|
|
description="The timestamp of the commit. Specify in YYYY-MM-DDTHH::MM:SSZ format",
|
|
type="datetime.datetime",
|
|
),
|
|
AttributeInfo(
|
|
name="author_name",
|
|
description="The name of the author of the commit",
|
|
type="string",
|
|
),
|
|
AttributeInfo(
|
|
name="author_email",
|
|
description="The email address of the author of the commit",
|
|
type="string",
|
|
),
|
|
]
|
|
# END DATASET SPECIFIC CODE
|
|
|
|
embeddings = OpenAIEmbeddings()
|
|
vectorstore = TimescaleVector(
|
|
embedding=embeddings,
|
|
collection_name=collection_name,
|
|
service_url=SERVICE_URL,
|
|
time_partition_interval=partition_interval,
|
|
)
|
|
|
|
llm = OpenAI(temperature=0)
|
|
retriever = SelfQueryRetriever.from_llm(
|
|
llm,
|
|
vectorstore,
|
|
document_content_description,
|
|
metadata_field_info,
|
|
enable_limit=True,
|
|
verbose=True,
|
|
)
|
|
|
|
template = """Answer the question based only on the following context:
|
|
{context}
|
|
|
|
Question: {question}
|
|
"""
|
|
prompt = ChatPromptTemplate.from_template(template)
|
|
|
|
model = ChatOpenAI(temperature=0, model="gpt-4")
|
|
|
|
# RAG chain
|
|
chain = (
|
|
RunnableParallel({"context": retriever, "question": RunnablePassthrough()})
|
|
| prompt
|
|
| model
|
|
| StrOutputParser()
|
|
)
|
|
|
|
|
|
class Question(BaseModel):
|
|
__root__: str
|
|
|
|
|
|
chain = chain.with_types(input_type=Question)
|