cohere

app.py

@@ -10,8 +10,8 @@ def get_emdedding_model():
     return Embeddings()
 
 
-def get_llm(url, api_key):
-    return LLM(url, api_key)
+def get_llm(api_key):
+    return LLM(api_key)
 
 
 def get_metadata(path):
@@ -23,7 +23,6 @@ def get_metadata(path):
             texts.append(data['text'])
     return texts, titles
 
-
 def combine_docs(indexes, texts):
     result = ""
     for i, index in enumerate(indexes):
@@ -32,64 +31,69 @@ def combine_docs(indexes, texts):
 
 
 def create_prompt(query, docs):
-    system_prompt = f""" You are a language model integrated into a search and
-    generation system based on relevant documents (RAG system).
-    Your task is to provide answers to the user's queries based on the provided
-    documents. Respond only based on the provided documents. Do not make up
-    information that is not in the sources. If you use data from a document,
-    indicate the document number in square brackets. For example: "This term
-    means such-and-such [1]." If there is no information in the documents,
-    politely explain that the information is not available. Do not alter the
-    content of the sources, convey the information accurately
-    Structure the text in a clear way whenever possible, even if formatting is
-    limited.
-    For example: 
-    User query: ML.
-    Documents:  
-    [1] es of ML models.  
-    [2] The rapid escalation of applying Machine Learning (ML) in various domains has led to paying more attention to the quality of ML components. There is then a growth of techniques and tools aiming at improving the quality of ML components and integrating them.  
-
-    Machine Learning (ML) is increasingly applied across various domains, leading to a focus on the quality of ML components and the development of techniques to improve and integrate them [2]
-
-    Follow this format in your responses and print all documents. User query: {query}. Documents: {docs}
+    system_prompt = f"""You are a language model integrated into a search and generation system based on relevant documents (RAG system).
+    Your task is to provide answers to the user's queries based solely on the provided documents.
+    If the information required to answer the user's question is available in the documents, use it, and refer to the document from which it was sourced by indicating its number in square brackets. For example: 
+    "This term means such-and-such [1]."
+    Ensure that the citation clearly refers to the relevant document and is placed directly after the information from the source.
+
+    If the information is not present in the documents, kindly explain that the information is not available, and do not speculate or make up information.
+
+    Do not alter the content or meaning of the sources. Convey the information accurately and structure your response clearly, even if the formatting options are limited.
+
+    User query: {query}
+    Documents:
+    {docs}
     """
     return system_prompt
 
 
-st.title("PaperRAG")
-st.write("RAG system for scientific papers with selectable search types")
+def main(query, search_types, llm_api_key):
+    model, llm = get_emdedding_model(), get_llm(llm_api_key)
+    texts, titles = get_metadata(config.PATH_METADATA)
+    embedding = model.get_query_embedding(query)
+
+    knowledge_base = KnowledgeBase(config.PATH_FAISS, config.PATH_PREPROCESSING_TEXT)
+    vector_search = []
+    bm25_search = []
 
-query = st.text_input("Enter your query", "")
-search_types = st.multiselect(
-    "Search Types", options=["Vector", "BM25"], default=["Vector", "BM25"]
-)
-llm_url = st.text_input("LLM URL", "", placeholder="Enter LLM ENDPOINT")
-llm_api_key = st.text_input("LLM API Key", "", placeholder="Enter LLM API Key", type="password")
+    if "Vector" in search_types:
+        vector_search = knowledge_base.search_by_embedding(embedding, 5)[0].tolist()
+    if "BM25" in search_types:
+        bm25_search = knowledge_base.search_by_BM25(query, 5)
 
-if st.button("Search"):
-    if query and llm_url and llm_api_key:
-        model = get_emdedding_model()
-        llm = get_llm(llm_url, llm_api_key)
+    docs = combine_docs(vector_search + bm25_search, texts)
+    prompt = create_prompt(query, docs)
 
-        texts, titles = get_metadata(config.PATH_METADATA)
-        embedding = model.get_query_embedding(query)
+    response = llm.generate_response(prompt)
+    return response, docs
 
-        knowledge_base = KnowledgeBase(config.PATH_FAISS, config.PATH_PREPROCESSING_TEXT)
 
-        vector_search = []
-        bm25_search = []
+# Streamlit Interface
+if __name__ == '__main__':
+    st.title("PaperRAG")
+    st.subheader("RAG system for scientific papers with selectable search types")
 
-        if "Vector" in search_types:
-            vector_search = knowledge_base.search_by_embedding(embedding, 5)[0].tolist()
-        if "BM25" in search_types:
-            bm25_search = knowledge_base.search_by_BM25(query, 5)
+    # User inputs
+    query = st.text_input("Enter your query")
+    search_types = st.multiselect(
+        "Select search types",
+        options=["Vector", "BM25"],
+        default=["Vector", "BM25"]
+    )
+    llm_api_key = st.text_input("Cohere API Key", type="password")
 
-        docs = combine_docs(vector_search + bm25_search, texts)
-        prompt = create_prompt(query, docs)
+    if st.button("Get Response"):
+        if query and llm_api_key:
+            # Call the main function
+            response, docs = main(query, search_types, llm_api_key)
 
-        response = llm.generate_response(prompt)
+            # Show the LLM response
+            st.subheader("LLM Response:")
+            st.text_area("Response", value=response, height=300)
 
-        st.subheader("Response")
-        st.write(response)
-    else:
-        st.error("Please fill in all the required fields.")
+            # Show combined documents
+            st.subheader("Citations:")
+            st.text_area("Documents", value=docs, height=300)
+        else:
+            st.error("Please enter both a query and an API key.")

requirements.txt

@@ -110,3 +110,4 @@ uvicorn==0.32.1
 websockets==12.0
 xxhash==3.5.0
 yarl==1.17.2
+litellm~=1.52.16

utils/llm.py

@@ -1,28 +1,15 @@
-import requests
-from dotenv import load_dotenv
-
-load_dotenv()
+import os
+from litellm import completion
 
 class LLM:
-    def __init__(self, url, api_key):
-        self.endpoint = url
-        self.api_key = api_key
-
-    def generate_response(self, prompt):
-        headers = {
-            "Content-Type": "application/json",
-            "api-key": self.api_key,
-        }
-
-        data = {
-            "messages": [{"role": "user", "content": prompt}],
-            "max_tokens": 3000,
-            "temperature": 0.5,
-        }
-
-        response = requests.post(self.endpoint, headers=headers, json=data)
-
-        if response.status_code == 200:
-            return response.json()["choices"][0]["message"]["content"]
-        else:
-            return ValueError(response.text)
+    def __init__(self, api_key):
+        os.environ["COHERE_API_KEY"] = api_key
+
+    def generate_response(self, prompt, temperature=0.5, max_tokens=1500):
+        response = completion(
+            model="command-r",
+            messages=[{"content": prompt, "role": "user"}],
+            temperature=temperature,
+            max_tokens=max_tokens
+        )
+        return response.choices[0].message.content