Update app.py

app.py

@@ -4,9 +4,10 @@ import gradio as gr
 from dotenv import load_dotenv
 from openai import AzureOpenAI
 
-from langchain_openai import AzureOpenAIEmbeddings
-
-from langchain_chroma import Chroma
+from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI
+from langchain.chains.query_constructor.base import AttributeInfo
+from langchain.retrievers.self_query.base import SelfQueryRetriever
+from langchain_community.vectorstores import Chroma
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain.retrievers.document_compressors import CrossEncoderReranker
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
@@ -16,7 +17,15 @@ load_dotenv()
 client = AzureOpenAI(
     api_key=os.environ['AZURE_OPENAI_KEY'],
     azure_endpoint=os.environ['AZURE_OPENAI_ENDPOINT'],
-    api_version='2024-02-01'
+    api_version=os.environ["AZURE_OPENAI_APIVERSION"]
+)
+
+llm = AzureChatOpenAI(
+    api_key=os.environ['AZURE_OPENAI_KEY'],
+    azure_endpoint=os.environ['AZURE_OPENAI_ENDPOINT'],
+    api_version=os.environ["AZURE_OPENAI_APIVERSION"],
+    model="gpt-4o-mini",
+    temperature=0
 )
 
 model_name = 'gpt-4o-mini'
@@ -24,7 +33,7 @@ model_name = 'gpt-4o-mini'
 embedding_model = AzureOpenAIEmbeddings(
     api_key=os.environ['AZURE_OPENAI_KEY'],
     azure_endpoint=os.environ['AZURE_OPENAI_ENDPOINT'],
-    api_version='2024-02-01',
+    api_version=os.environ["AZURE_OPENAI_APIVERSION"],
     azure_deployment="text-embedding-ada-002"
 )
 
@@ -36,8 +45,38 @@ vectorstore_persisted = Chroma(
     embedding_function=embedding_model
 )
 
-retriever = vectorstore_persisted.as_retriever(
-    search_type='similarity',
+metadata_field_info = [
+    AttributeInfo(
+        name="year",
+        description="The year of the Tesla 10-K annual report",
+        type="string",
+    ),
+    AttributeInfo(
+        name="file",
+        description="The filename of the source document",
+        type="string",
+    ),
+    AttributeInfo(
+        name="page_number",
+        description="The page number of the document in the original file",
+        type="string",
+    ),
+    AttributeInfo(
+        name="source", 
+        description="The source of the document content: text or image", 
+        type="string"
+    )
+]
+
+document_content_description = "10-k Statements from Tesla"
+
+retriever = SelfQueryRetriever.from_llm(
+    llm,
+    vectorstore_persisted,
+    document_content_description,
+    metadata_field_info,
+    enable_limit=True,
+    verbose=True,
     search_kwargs={'k': 10}
 )