Update app.py

app.py

@@ -1,257 +1,249 @@
 import os
-import logging
-import asyncio
-import random  # Import random for token selection
-from typing import AsyncGenerator, Tuple
+import json
+import re
 import gradio as gr
-from huggingface_hub import InferenceClient
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
-from langchain.schema import Document
+import requests
 from duckduckgo_search import DDGS
+from typing import List
+from pydantic import BaseModel, Field
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_core.documents import Document
+from huggingface_hub import InferenceClient
+import logging
 
-# Configure logging
+# Set up basic configuration for logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# List of Hugging Face tokens
-huggingface_tokens = [
-    os.environ.get("HUGGINGFACE_TOKEN_1"),
-    os.environ.get("HUGGINGFACE_TOKEN_2"),
-    os.environ.get("HUGGINGFACE_TOKEN_3")
-]
-
-# Function to get a random Hugging Face token
-def get_random_token():
-    return random.choice(huggingface_tokens)
+# Environment variables and configurations
+huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
+ACCOUNT_ID = os.environ.get("CLOUDFARE_ACCOUNT_ID")
+API_TOKEN = os.environ.get("CLOUDFLARE_AUTH_TOKEN")
+API_BASE_URL = "https://api.cloudflare.com/client/v4/accounts/a17f03e0f049ccae0c15cdcf3b9737ce/ai/run/"
 
 MODELS = [
     "mistralai/Mistral-7B-Instruct-v0.3",
     "mistralai/Mixtral-8x7B-Instruct-v0.1",
-    "mistralai/Mistral-Nemo-Instruct-2407",
-    "meta-llama/Meta-Llama-3.1-8B-Instruct",
-    "meta-llama/Meta-Llama-3.1-70B-Instruct",
-    "google/gemma-2-9b-it",
-    "google/gemma-2-27b-it"
+    "@cf/meta/llama-3.1-8b-instruct",
+    "mistralai/Mistral-Nemo-Instruct-2407"
 ]
 
-DEFAULT_SYSTEM_PROMPT = """You are a world-class financial AI assistant, capable of complex reasoning and reflection. 
-Reason through the query inside <thinking> tags, and then provide your final response inside <output> tags.
-Providing comprehensive and accurate information based on web search results is essential. 
-Your goal is to synthesize the given context into a coherent and detailed response that directly addresses the user's query. 
-Please ensure that your response is well-structured and factual.
-If you detect that you made a mistake in your reasoning at any point, correct yourself inside <reflection> tags."""
-
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
 def duckduckgo_search(query):
+    with DDGS() as ddgs:
+        results = ddgs.text(query, max_results=5)
+    return results
+
+class CitingSources(BaseModel):
+    sources: List[str] = Field(
+        ...,
+        description="List of sources to cite. Should be an URL of the source."
+    )
+
+def chatbot_interface(message, history, model, temperature, num_calls):
+    if not message.strip():
+        return "", history
+
+    history = history + [(message, "")]
+
     try:
-        with DDGS() as ddgs:
-            results = ddgs.text(query, max_results=5)
-        logging.info(f"Search completed for query: {query}")
-        return results
+        for response in respond(message, history, model, temperature, num_calls):
+            history[-1] = (message, response)
+            yield history
+    except gr.CancelledError:
+        yield history
     except Exception as e:
-        logging.error(f"Error during DuckDuckGo search: {str(e)}")
-        return []
+        logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
+        history[-1] = (message, f"An unexpected error occurred: {str(e)}")
+        yield history
+
+def retry_last_response(history, model, temperature, num_calls):
+    if not history:
+        return history
+    
+    last_user_msg = history[-1][0]
+    history = history[:-1]  # Remove the last response
+    
+    return chatbot_interface(last_user_msg, history, model, temperature, num_calls)
+
+def respond(message, history, model, temperature, num_calls):
+    logging.info(f"User Query: {message}")
+    logging.info(f"Model Used: {model}")
+
+    try:
+        for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+            response = f"{main_content}\n\n{sources}"
+            first_line = response.split('\n')[0] if response else ''
+            yield response
+    except Exception as e:
+        logging.error(f"Error with {model}: {str(e)}")
+        yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
 def create_web_search_vectors(search_results):
     embed = get_embeddings()
+    
     documents = []
     for result in search_results:
         if 'body' in result:
             content = f"{result['title']}\n{result['body']}\nSource: {result['href']}"
             documents.append(Document(page_content=content, metadata={"source": result['href']}))
-    logging.info(f"Created vectors for {len(documents)} search results.")
+    
     return FAISS.from_documents(documents, embed)
 
-def create_context(search_results, use_embeddings, query):
-    if use_embeddings:
-        web_search_database = create_web_search_vectors(search_results)
-        retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
-        relevant_docs = retriever.get_relevant_documents(query)
-        return "\n".join([doc.page_content for doc in relevant_docs])
-    else:
-        return "\n".join([f"{result['title']}\n{result['body']}" for result in search_results])
-
-async def get_response_with_search(query: str, system_prompt: str, model: str, use_embeddings: bool, history=None, num_calls: int = 3, temperature: float = 0.2) -> AsyncGenerator[Tuple[str, str], None]:
+def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     search_results = duckduckgo_search(query)
+    web_search_database = create_web_search_vectors(search_results)
     
-    if not search_results:
-        logging.warning(f"No web search results found for query: {query}")
+    if not web_search_database:
         yield "No web search results available. Please try again.", ""
         return
-
-    sources = [result['href'] for result in search_results if 'href' in result]
-    source_list_str = "\n".join(sources)
-
-    context = create_context(search_results, use_embeddings, query)
-    logging.info(f"Context created for query: {query}")
-
-    user_message = f"""Using the following context from web search results:
+    
+    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
+    relevant_docs = retriever.get_relevant_documents(query)
+    
+    context = "\n".join([doc.page_content for doc in relevant_docs])
+    
+    prompt = f"""Using the following context from web search results:
 {context}
+Write a detailed and complete research document that fulfills the following user request: '{query}'
+After writing the document, please provide a list of sources used in your response."""
 
-Write a detailed and complete research document that fulfills the following user request: '{query}'."""
+    if model == "@cf/meta/llama-3.1-8b-instruct":
+        # Use Cloudflare API
+        for response in get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature):
+            yield response, ""  # Yield streaming response without sources
+    else:
+        # Use Hugging Face API
+        client = InferenceClient(model, token=huggingface_token)
+        
+        main_content = ""
+        for i in range(num_calls):
+            for message in client.chat_completion(
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=10000,
+                temperature=temperature,
+                stream=True,
+            ):
+                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                    chunk = message.choices[0].delta.content
+                    main_content += chunk
+                    yield main_content, ""  # Yield partial main content without sources
+
+def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2):
+    headers = {
+        "Authorization": f"Bearer {API_TOKEN}",
+        "Content-Type": "application/json"
+    }
+    model = "@cf/meta/llama-3.1-8b-instruct"
+
+    instruction = f"""Using the following context:
+{context}
+Write a detailed and complete research document that fulfills the following user request: '{query}'
+After writing the document, please provide a list of sources used in your response."""
 
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": user_message}
+    inputs = [
+        {"role": "system", "content": instruction},
+        {"role": "user", "content": query}
     ]
 
-    if history:
-        messages = history + messages
+    payload = {
+        "messages": inputs,
+        "stream": True,
+        "temperature": temperature,
+        "max_tokens": 32000
+    }
 
-    # Get a random token for the API call
-    token = get_random_token()
-    client = InferenceClient(model, token=token)
     full_response = ""
-
-    for call in range(num_calls):
+    for i in range(num_calls):
         try:
-            response = await asyncio.to_thread(
-                client.chat_completion,
-                messages=messages,
-                max_tokens=6000,
-                temperature=temperature,
-                top_p=0.8,
-            )
-
-            if response is None or not isinstance(response, dict) or 'choices' not in response:
-                logging.error(f"API call {call + 1} returned an invalid response: {response}")
-                if call == num_calls - 1:
-                    yield "The API returned an invalid response. Please try again later.", ""
-                continue
-
-            new_content = response['choices'][0]['message']['content']
-            full_response += new_content
-            yield full_response, ""
-
-            if full_response:
-                break  # If we got a valid response, exit the loop
-
+            with requests.post(f"{API_BASE_URL}{model}", headers=headers, json=payload, stream=True) as response:
+                if response.status_code == 200:
+                    for line in response.iter_lines():
+                        if line:
+                            try:
+                                json_response = json.loads(line.decode('utf-8').split('data: ')[1])
+                                if 'response' in json_response:
+                                    chunk = json_response['response']
+                                    full_response += chunk
+                                    yield full_response
+                            except (json.JSONDecodeError, IndexError) as e:
+                                logging.error(f"Error parsing streaming response: {str(e)}")
+                                continue
+                else:
+                    logging.error(f"HTTP Error: {response.status_code}, Response: {response.text}")
+                    yield f"I apologize, but I encountered an HTTP error: {response.status_code}. Please try again later."
         except Exception as e:
-            logging.error(f"Error in API call {call + 1}: {str(e)}")
-            if call == num_calls - 1:
-                yield f"An error occurred during API calls: {str(e)}. Please try again later.", ""
-        
-        await asyncio.sleep(1)  # 1 second delay between calls
-
+            logging.error(f"Error in generating response from Cloudflare: {str(e)}")
+            yield f"I apologize, but an error occurred: {str(e)}. Please try again later."
+    
     if not full_response:
-        logging.warning("No response generated from the model")
-        yield "No response generated from the model. Please try again.", ""
-    else:
-        yield f"{full_response}\n\nSources:\n{source_list_str}", ""
-
-def process_history(history):
-    chat_history = []
-    if isinstance(history, str):
-        # If history is a string (like the system prompt), add it as a system message
-        chat_history.append({"role": "system", "content": history})
-    elif isinstance(history, list):
-        for entry in history:
-            if isinstance(entry, (list, tuple)) and len(entry) == 2:
-                human, assistant = entry
-                chat_history.append({"role": "user", "content": human})
-                if assistant:
-                    chat_history.append({"role": "assistant", "content": assistant})
-            elif isinstance(entry, str):
-                # If it's a string, assume it's a user message
-                chat_history.append({"role": "user", "content": entry})
-    return chat_history
-
-async def respond(message, system_prompt, history, model, temperature, num_calls, use_embeddings):
-    logging.info(f"User Query: {message}")
-    logging.info(f"Model Used: {model}")
-    logging.info(f"Temperature: {temperature}")
-    logging.info(f"Number of API Calls: {num_calls}")
-    logging.info(f"Use Embeddings: {use_embeddings}")
-    logging.info(f"System Prompt: {system_prompt}")
-    logging.info(f"History: {history}")
-
-    chat_history = process_history(history)
-
-    try:
-        async for main_content, sources in get_response_with_search(
-            message, 
-            system_prompt, 
-            model, 
-            use_embeddings, 
-            history=chat_history, 
-            num_calls=num_calls, 
-            temperature=temperature
-        ):
-            yield main_content
-
-        if sources:
-            yield f"\n\nSources:\n{sources}"
+        yield "I apologize, but I couldn't generate a response at this time. Please try again later."
 
-    except asyncio.CancelledError:
-        logging.warning("The operation was cancelled.")
-        yield "The operation was cancelled. Please try again."
-    except Exception as e:
-        logging.error(f"Error in respond function: {str(e)}")
-        yield f"An error occurred: {str(e)}"
+def vote(data: gr.LikeData):
+    if data.liked:
+        print(f"You upvoted this response: {data.value}")
+    else:
+        print(f"You downvoted this response: {data.value}")
 
 css = """
 /* Fine-tune chatbox size */
-.chatbot-container {
-    height: 600px !important;
-    width: 100% !important;
-}
-.chatbot-container > div {
-    height: 100%;
-    width: 100%;
-}
 """
 
-# Gradio interface setup
-def create_gradio_interface():
-    custom_placeholder = "Enter your question here for web search."
+def initial_conversation():
+    return [
+        (None, "Welcome! I'm your AI assistant for web search. Here's how you can use me:\n\n"
+                "1. Ask me any question, and I'll search the web for information.\n"
+                "2. You can adjust the model, temperature, and number of API calls for fine-tuned responses.\n"
+                "3. For any queries, feel free to reach out @[email protected] or discord - shreyas094\n\n"
+                "To get started, ask me a question!")
+    ]
 
-    demo = gr.ChatInterface(
-        fn=respond,
-        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=True, render=False),
-        additional_inputs=[
-            gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, lines=6, label="System Prompt", placeholder="Enter your system prompt here"),
-            gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[3]),
-            gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
-            gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
-            gr.Checkbox(label="Use Embeddings", value=False),
-        ],
-        title="AI-powered Web Search Assistant",
-        description="Use web search to answer questions or generate summaries.",
-        theme=gr.Theme.from_hub("allenai/gradio-theme"),
-        css=css,
-        examples=[
-            ["What are the latest developments in artificial intelligence?"],
-            ["Explain the concept of quantum computing."],
-            ["What are the environmental impacts of renewable energy?"]
-        ],
-        cache_examples=False,
-        analytics_enabled=False,
-        textbox=gr.Textbox(placeholder=custom_placeholder, container=False, scale=7),
-        chatbot=gr.Chatbot(
-            show_copy_button=True,
-            likeable=True,
-            layout="bubble",
-            height=400,
-        )
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[3]),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
+        gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
+    ],
+    title="AI-powered Web Search Assistant",
+    description="Ask questions and get answers from web search results.",
+    theme=gr.themes.Soft(
+        primary_hue="orange",
+        secondary_hue="amber",
+        neutral_hue="gray",
+        font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]
+    ).set(
+        body_background_fill_dark="#0c0505",
+        block_background_fill_dark="#0c0505",
+        block_border_width="1px",
+        block_title_background_fill_dark="#1b0f0f",
+        input_background_fill_dark="#140b0b",
+        button_secondary_background_fill_dark="#140b0b",
+        border_color_accent_dark="#1b0f0f",
+        border_color_primary_dark="#1b0f0f",
+        background_fill_secondary_dark="#0c0505",
+        color_accent_soft_dark="transparent",
+        code_background_fill_dark="#140b0b"
+    ),
+    css=css,
+    examples=[
+        ["What are the latest developments in artificial intelligence?"],
+        ["Can you explain the basics of quantum computing?"],
+        ["What are the current global economic trends?"]
+    ],
+    cache_examples=False,
+    analytics_enabled=False,
+    textbox=gr.Textbox(placeholder="Ask a question", container=False, scale=7),
+    chatbot = gr.Chatbot(  
+        show_copy_button=True,
+        likeable=True,
+        layout="bubble",
+        height=400,
+        value=initial_conversation()
     )
-
-    with demo:
-        gr.Markdown("""
-        ## How to use
-        1. Enter your question in the chat interface.
-        2. Optionally, modify the System Prompt to guide the AI's behavior.
-        3. Select the model you want to use from the dropdown.
-        4. Adjust the Temperature to control the randomness of the response.
-        5. Set the Number of API Calls to determine how many times the model will be queried.
-        6. Check or uncheck the "Use Embeddings" box to toggle between using embeddings or direct text summarization.
-        7. Press Enter or click the submit button to get your answer.
-        8. Use the provided examples or ask your own questions.
-        """)
-
-    return demo
+)
 
 if __name__ == "__main__":
-    demo = create_gradio_interface()
     demo.launch(share=True)