Upload app.py

app.py

@@ -59,11 +59,16 @@ def save_key(api_key):
     return api_key
 
 
-def query_pinecone(query, top_k, model, index, threshold=0.5):
+def query_pinecone(query, top_k, model, index, year, quarter, ticker, threshold=0.5):
     # generate embeddings for the query
     xq = model.encode([query]).tolist()
     # search pinecone index for context passage with the answer
-    xc = index.query(xq, top_k=top_k, include_metadata=True)
+    xc = index.query(
+        xq,
+        top_k=top_k,
+        filter={"year": year, "quarter": quarter, "ticker": ticker},
+        include_metadata=True,
+    )
     # filter the context passages based on the score threshold
     filtered_matches = []
     for match in xc["matches"]:
@@ -137,6 +142,27 @@ st.write(
 
 query_text = st.text_input("Input Query", value="Who is the CEO of Apple?")
 
+years_choice = ["2016", "2017", "2018", "2019", "2020"]
+
+year = st.selectbox("Year", years_choice)
+
+quarter = st.selectbox("Quarter", ["Q1", "Q2", "Q3", "Q4"])
+
+ticker_choice = [
+    "AAPL",
+    "CSCO",
+    "MSFT",
+    "ASML",
+    "NVDA",
+    "GOOGL",
+    "MU",
+    "INTC",
+    "AMZN",
+    "AMD",
+]
+
+ticker = st.selectbox("Company", ticker_choice)
+
 num_results = int(st.number_input("Number of Results to query", 1, 5, value=3))
 
 
@@ -180,7 +206,14 @@ threshold = float(
 data = get_data()
 
 query_results = query_pinecone(
-    query_text, num_results, retriever_model, pinecone_index, threshold
+    query_text,
+    num_results,
+    retriever_model,
+    pinecone_index,
+    year,
+    quarter,
+    ticker,
+    threshold,
 )
 
 if threshold <= 0.60: