app.py

import os
import json
import logging
import gradio as gr
from duckduckgo_search import DDGS
from typing import List, Dict
from pydantic import BaseModel
from huggingface_hub import InferenceClient
from typing import List, Dict
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# Set up basic configuration for logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Environment variables and configurations
huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")

MODELS = [
    "mistralai/Mistral-7B-Instruct-v0.3",
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "duckduckgo/gpt-4o-mini",
    "duckduckgo/claude-3-haiku",
    "duckduckgo/llama-3.1-70b",
    "duckduckgo/mixtral-8x7b"
]

class ConversationManager:
    def __init__(self):
        self.history = []
        self.current_context = None

    def add_interaction(self, query, response):
        self.history.append((query, response))
        self.current_context = f"Previous query: {query}\nPrevious response summary: {response[:200]}..."

    def get_context(self):
        return self.current_context

conversation_manager = ConversationManager()

huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")

def duckduckgo_search(query: str, max_results: int = 10) -> List[Dict[str, str]]:
    with DDGS() as ddgs:
        results = list(ddgs.text(query, max_results=max_results))
    return results

def get_web_search_results(query: str, model: str, num_calls: int = 3, temperature: float = 0.2, max_results: int = 10, search_method: str = "DDGS.chat") -> Dict[str, any]:
    try:
        # Perform web search
        if search_method == "DDGS.text":
            search_results = duckduckgo_search(query, max_results)
        else:  # Default to DDGS.chat
            with DDGS() as ddgs:
                search_results = list(ddgs.text(query, max_results=max_results))
        
        if not search_results:
            return {"error": f"No results found for query: {query}"}
        
        # Create embeddings for search results
        embedder = SentenceTransformer('all-MiniLM-L6-v2')
        web_search_vectors = embedder.encode([result['body'] for result in search_results])
        
        # Retrieve relevant documents
        query_vector = embedder.encode([query])
        similarities = cosine_similarity(query_vector, web_search_vectors)[0]
        top_indices = np.argsort(similarities)[-5:][::-1]
        relevant_docs = [search_results[i] for i in top_indices]
        
        # Prepare context
        context = "\n".join([f"Title: {doc['title']}\nContent: {doc['body']}" for doc in relevant_docs])
        
        # Prepare prompt
        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."""
        
        # Generate response based on the selected model
        if model == "@cf/meta/llama-3.1-8b-instruct":
            # Use Cloudflare API (placeholder, as the actual implementation is not provided)
            response = get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature, search_type="web")
        else:
            # Use Hugging Face API
            client = InferenceClient(model, token=huggingface_token)
            response = ""
            for _ 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:
                        response += message.choices[0].delta.content
        
        return {
            "query": query,
            "search_results": search_results,
            "relevant_docs": relevant_docs,
            "response": response
        }
    
    except Exception as e:
        return {"error": f"An error occurred during web search and processing: {str(e)}"}


def rephrase_query(original_query: str, conversation_manager: ConversationManager) -> str:
    context = conversation_manager.get_context()
    if context:
        prompt = f"""You are a highly intelligent conversational chatbot. Your task is to analyze the given context and new query, then decide whether to rephrase the query with or without incorporating the context. Follow these steps:

        1. Determine if the new query is a continuation of the previous conversation or an entirely new topic.
        2. If it's a continuation, rephrase the query by incorporating relevant information from the context to make it more specific and contextual.
        3. If it's a new topic, rephrase the query to make it more appropriate for a web search, focusing on clarity and accuracy without using the previous context.
        4. Provide ONLY the rephrased query without any additional explanation or reasoning.
        
        Context: {context}
        
        New query: {original_query}
        
        Rephrased query:"""
        response = DDGS().chat(prompt, model="llama-3.1-70b")
        rephrased_query = response.split('\n')[0].strip()
        return rephrased_query
    return original_query

def summarize_web_results(query: str, search_results: List[Dict[str, str]], conversation_manager: ConversationManager) -> str:
    try:
        context = conversation_manager.get_context()
        search_context = "\n\n".join([f"Title: {result['title']}\nContent: {result['body']}" for result in search_results])

        prompt = f"""You are a highly intelligent & expert analyst and your job is to skillfully articulate the web search results about '{query}' and considering the context: {context}, 
        You have to create a comprehensive news summary FOCUSING on the context provided to you. 
        Include key facts, relevant statistics, and expert opinions if available. 
        Ensure the article is well-structured with an introduction, main body, and conclusion, IF NECESSARY. 
        Address the query in the context of the ongoing conversation IF APPLICABLE.
        Cite sources directly within the generated text and not at the end of the generated text, integrating URLs where appropriate to support the information provided:

        {search_context}

        Article:"""

        summary = DDGS().chat(prompt, model="llama-3.1-70b")
        return summary
    except Exception as e:
        return f"An error occurred during summarization: {str(e)}"

def respond(message, history, model, temperature, num_calls, use_web_search, search_method):
    logging.info(f"User Query: {message}")
    logging.info(f"Model Used: {model}")
    logging.info(f"Use Web Search: {use_web_search}")
    logging.info(f"Search Method: {search_method}")

    if use_web_search:
        original_query = message
        rephrased_query = rephrase_query(message, conversation_manager)
        logging.info(f"Original query: {original_query}")
        logging.info(f"Rephrased query: {rephrased_query}")

        final_summary = ""
        for _ in range(num_calls):
            search_results = get_web_search_results(rephrased_query, model, num_calls, temperature, max_results=10, search_method=search_method)
            if not search_results:
                final_summary += f"No search results found for the query: {rephrased_query}\n\n"
            elif "error" in search_results[0]:
                final_summary += search_results[0]["error"] + "\n\n"
            else:
                summary = summarize_web_results(rephrased_query, search_results, conversation_manager)
                final_summary += summary + "\n\n"

        if final_summary:
            conversation_manager.add_interaction(original_query, final_summary)
            yield final_summary
        else:
            yield "Unable to generate a response. Please try a different query."
    else:
        yield "Web search is not enabled. Please enable web search to use this feature."

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%;
}
"""

def initial_conversation():
    return [
        (None, "Welcome! I'm your AI assistant for web search. Here's how you can use me:\n\n"
                "1. Make sure the 'Use Web Search' checkbox is enabled in the Additional Inputs section.\n"
                "2. Ask me any question, and I'll search the web for the most relevant and up-to-date information.\n"
                "3. You can adjust the model, temperature, number of API calls, and search method in the Additional Inputs section to fine-tune your results.\n\n"
                "To get started, just ask me a question!")
    ]

# Create the Gradio interface
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
        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 Web Search", value=True),
        gr.Radio(["DDGS.chat", "DDGS.text"], label="Search Method", value="DDGS.chat")
    ],
    title="AI-powered Web Search Assistant",
    description="Ask questions and get answers from the latest web information.",
    theme=gr.Theme.from_hub("allenai/gradio-theme"),
    css=css,
    examples=[
        ["What's the latest news about artificial intelligence?"],
        ["Summarize the current global economic situation."],
        ["What are the top environmental concerns right now?"],
        ["What are the recent breakthroughs in quantum computing?"]
    ],
    cache_examples=False,
    analytics_enabled=False,
    textbox=gr.Textbox(placeholder="Ask any question", container=False, scale=7),
    chatbot = gr.Chatbot(  
        show_copy_button=True,
        likeable=True,
        layout="bubble",
        height=400,
        value=initial_conversation()
    )
)

if __name__ == "__main__":
    demo.launch(share=True)