Upload retriever.py

retriever.py

@@ -0,0 +1,53 @@
+from google import genai
+import time
+
+# set up authentication
+api_key = userdata.get("GEMINI_API_KEY")
+client = genai.Client(api_key=api_key)
+
+# read in documents
+acled = PdfReader("/content/sample_data/ACLED.pdf")
+dame = PdfReader("/content/sample_data/DAME.pdf")
+nansen = PdfReader("/content/sample_data/Nansen.pdf")
+
+# extract text from documents
+acled_text = ""
+
+dame_text = ""
+
+nansen_text = ""
+
+for page in acled.pages:
+    acled_text += page.extract_text()
+
+for page in dame.pages:
+    dame_text += page.extract_text()
+
+for page in nansen.pages:
+    nansen_text += page.extract_text()
+
+# place extracted texts into a single list
+documents = [acled_text, dame_text, nansen_text]
+
+def retriever(prompt, history):
+  """
+  Retrieves information from pre-loaded documents.
+  Args:
+      prompt: A string representing the search query
+      history: A placeholder representing query history
+
+  Returns:
+      Search results in natural language.
+  """
+  context = "\n\n".join(documents)
+  super_prompt = f"{context}\n\nRetrieve information:\n\nQuestion: {prompt}"
+
+  response = client.models.generate_content(
+    model="gemini-2.0-flash",
+    contents=super_prompt)
+
+  stream_text = ""
+  for each in response.text.split():
+    stream_text += each + " "
+    time.sleep(0.01)
+    yield stream_text