add main

main.py

@@ -0,0 +1,184 @@
+import gradio as gr
+import os
+from langchain.callbacks.base import BaseCallbackHandler
+from collections.abc import Generator
+from queue import Queue, Empty
+from threading import Thread
+
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from langchain import PromptTemplate
+from langchain.chains import LLMChain
+from langchain.chat_models import ChatOpenAI
+import pinecone
+from langchain.embeddings import OpenAIEmbeddings
+
+
+OPENAI_API_KEY=os.environ["OPENAI_API_KEY"]
+PINECONE_API_KEY=os.environ["PINECONE_API_KEY"]
+PINECONE_ENV=os.environ["PINECONE_ENV"]
+PINECONE_INDEX=os.environ["PINECONE_INDEX"]
+
+class QueueCallback(BaseCallbackHandler):
+    """Callback handler for streaming LLM responses to a queue."""
+
+    def __init__(self, q):
+        self.q = q
+
+    def on_llm_new_token(self, token: str, **kwargs: any) -> None:
+        self.q.put(token)
+
+    def on_llm_end(self, *args, **kwargs: any) -> None:
+        return self.q.empty()
+
+# TOOL
+#####################################################################
+llm = ChatOpenAI(model_name="gpt-4-1106-preview", temperature=0)
+
+template = """
+You are an expert research assistant. You can access information about articles via your tool. 
+Use information ONLY from this tool. Do not invent or add any more knowladge, be strict for the articles.
+{instuction}
+
+User: {user}
+--------
+{content}
+"""
+
+prompt = PromptTemplate(
+    input_variables=["instuction", "user", "content"],
+    template=template,
+)
+
+chain = LLMChain(llm=llm, prompt=prompt, callbacks=[QueueCallback])
+
+pinecone.init(
+    api_key=PINECONE_API_KEY, 
+    environment=PINECONE_ENV
+)
+
+index = pinecone.Index(PINECONE_INDEX)
+embedder = OpenAIEmbeddings()
+
+
+class PineconeSearch:
+    docsearch
+    topk
+
+    def __init__(
+        namespace,
+        topk
+        ):
+        self.docsearch = Pinecone.from_existing_index(PINECONE_INDEX, embedder, namespace=namespace)
+        self.topk=topk
+
+    def __call__(query):
+        response = self.docsearch.similarity_search(query=query, k=self.topk)
+        context = ""
+        for doc in docs:
+            context += f"Coontent:\n{doc.page_content}\n"
+            context += f"Source: {doc.metadta['url']}\n"
+            contect += "----"
+        return context
+
+
+
+def query_tool(category, pinecone_topk, query):
+    data = {
+        "1_D3_receptor": "demo-richter-target-400-30-1",
+        "2_dopamine": "demo-richter-target-400-30-2",
+        "3_mitochondrial": "demo-richter-target-400-30-3"
+    }
+    
+    pinecone_namespace = data[category]
+
+    search_tool = PineconeSearch(
+        namespace=pinecone_namespace,
+        topk=pinecone_topk,
+    )
+
+    return search_tool(query)
+
+
+
+def print_token_and_price(response):
+    inp = sum(response["token_usage"]["prompt_tokens"])
+    out = sum( response["token_usage"]["completion_tokens"])
+    print(f"Token usage: {inp+out}")
+    price = inp/1000*0.01 + out/1000*0.03
+    print(f"Total price: {price*370:.2f} Ft")
+    print("===================================")
+
+
+
+def stream(input_text, history, user_prompt, topic, topk) -> Generator:
+    # Create a Queue
+    q = Queue()
+    job_done = object()
+
+    # Create a funciton to call - this will run in a thread
+    def task():
+        tool_resp = query_tool(topic, topk, input_text)
+
+        response = chain({"instuction": user_prompt, "user": input_text, "content": tool_resp})
+
+        #print_token_and_price(response=response)
+        q.put(job_done)
+
+    # Create a thread and start the function
+    t = Thread(target=task)
+    t.start()
+
+    content = ""
+
+    # Get each new token from the queue and yield for our generator
+    counter = 0
+    while True:
+        try:
+            next_token = q.get(True, timeout=1)
+            if next_token is job_done:
+                break
+            content += next_token
+            counter += 1
+            if counter == 20:
+                content += "\n"
+                counter = 0
+            if "\n" in next_token:
+                counter = 0
+            yield next_token, content
+        except Empty:
+            continue
+
+def ask_llm(message, history, prompt, topic, topk):
+    for next_token, content in stream(message, history, prompt, topic, topk):
+        yield(content)
+
+
+agent_prompt_textbox = gr.Textbox(
+    label = "Set the behaviour of the agent",
+    lines = 2,
+    value = "Make your brief answer in bullet points."
+)
+namespace_drobdown = gr.Dropdown(
+    ["1_D3_receptor", "2_dopamine", "3_mitochondrial"],
+    label="Choose a topic",
+    value="1_D3_receptor"
+    )
+topk_slider = gr.Slider(
+    minimum=10,
+    maximum=350,
+    value=70,
+    step=10
+)
+
+
+additional_inputs = [agent_prompt_textbox, namespace_drobdown, topk_slider]
+
+chatInterface = gr.ChatInterface(
+    fn=ask_llm,
+    additional_inputs=additional_inputs,
+    additional_inputs_accordion_name="Agent parameters"
+    ).queue().launch()