add openrouter api

app.py

@@ -1,23 +1,14 @@
-import urllib.request 
-import fitz
-import re
-import numpy as np
-import tensorflow_hub as hub
 from openai import OpenAI
 import gradio as gr
 import os
 import shutil
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from sklearn.neighbors import NearestNeighbors
-import anthropic
 
-# client = OpenAI(
-#     base_url='https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1/v1/',
-#     api_key=os.getenv('openai_key')
-# )
-
-client = anthropic.Anthropic()
+client = OpenAI(
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv('OPENROUTER_API_KEY')
+)
 
 from util import pdf_to_text, text_to_chunks, SemanticSearch
 
@@ -30,45 +21,14 @@ def load_recommender(path, start_page=1):
     return 'Corpus Loaded.'
 
 
-# def openai_generate_text(prompt, model = "gpt-3.5-turbo-16k-0613"):
-#     model="mistralai/Mixtral-8x7B-Instruct-v0.1"
-#     max_tokens=1024
-#     message = clinet.chat.completions.create(
-#         model=model,
-#         messages=[
-#             {"role": "user", "content": prompt}
-#         ],
-#         max_tokens=max_tokens,
-#     ).choices[0].message.content
-#     return message
-
-def claude_generate_text(prompt, model = "claude-3-haiku-20240307"):
-    message = client.messages.create(
-        model=model,
-        max_tokens=1000,
-        temperature=0.0,
-        # system="Respond only in mandarin",
+def generate_text(prompt):
+    message = client.chat.completions.create(
+        model="google/gemini-pro",
         messages=[
             {"role": "user", "content": prompt}
-        ]
-    )
-    return message.content[0].text
-
-def generate_answer(question):
-    topn_chunks = recommender(question)
-    prompt = 'search results:\n\n'
-    for c in topn_chunks:
-        prompt += c + '\n\n'
-        
-    prompt += "Instructions: Compose a comprehensive reply to the query using the search results given. "\
-              "Cite each reference using [ Page Number] notation. "\
-              "Only answer what is asked. The answer should be short and concise. "\
-              "If asked in Chinese, respond in Chinese; if asked in English, respond"\
-              "in English \n\nQuery: "
-    
-    prompt += f"{question}\nAnswer:"
-    answer = claude_generate_text(prompt)
-    return answer
+        ],
+    ).choices[0].message.content
+    return message
 
 def question_answer(chat_history, file, question):
     suffix = Path(file.name).suffix
@@ -77,7 +37,7 @@ def question_answer(chat_history, file, question):
         tmp_path = Path(tmp.name)
 
     load_recommender(str(tmp_path))
-    answer = generate_answer(question)
+    answer = generate_text(question)
     chat_history.append([question, answer])
     return chat_history