app.py

import os
import logging
import asyncio
import gradio as gr
from huggingface_hub import InferenceClient
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.schema import Document
from duckduckgo_search import DDGS

# Configure 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",
    "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"
]

# Default system message template
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):
    try:
        with DDGS() as ddgs:
            results = ddgs.text(query, max_results=5)
        logging.info(f"Search completed for query: {query}")
        return results
    except Exception as e:
        logging.error(f"Error during DuckDuckGo search: {str(e)}")
        return []

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)

async def get_response_with_search(query, system_prompt, model, use_embeddings, history=None, num_calls=3, temperature=0.2):
    search_results = duckduckgo_search(query)
    
    if not search_results:
        logging.warning(f"No web search results found for query: {query}")
        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)

    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)
        context = "\n".join([doc.page_content for doc in relevant_docs])
    else:
        context = "\n".join([f"{result['title']}\n{result['body']}" for result in search_results])

    logging.info(f"Context created for query: {query}")

    user_message = f"""Using the following context from web search results:
{context}

Write a detailed and complete research document that fulfills the following user request: '{query}'."""

    client = InferenceClient(model, token=huggingface_token)
    full_response = ""
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_message}
    ]

    # Include chat history if provided
    if history:
        messages = history + messages

    try:
        for call in range(num_calls):
            try:
                for response in client.chat_completion(
                    messages=messages,
                    max_tokens=6000,
                    temperature=temperature,
                    stream=True,
                    top_p=0.8,
                ):
                    if isinstance(response, dict) and "choices" in response:
                        for choice in response["choices"]:
                            if "delta" in choice and "content" in choice["delta"]:
                                chunk = choice["delta"]["content"]
                                full_response += chunk
                                yield full_response, ""
                    else:
                        logging.error("Unexpected response format or missing attributes in the response object.")
                        break
            except Exception as e:
                logging.error(f"Error in API call {call + 1}: {str(e)}")
                if "422 Client Error" in str(e):
                    logging.warning("Received 422 Client Error. Adjusting request parameters.")
                    # You might want to adjust parameters here, e.g., reduce max_tokens
                yield f"An error occurred during API call {call + 1}. Retrying...", ""
            
            # Add a small delay between API calls
            await asyncio.sleep(1)  # 1 second delay

    except asyncio.CancelledError:
        logging.warning("The operation was cancelled.")
        yield "The operation was cancelled. Please try again.", ""

    if not full_response:
        logging.warning("No response generated from the model")
        yield "No response generated from the model.", ""

    yield f"{full_response}\n\nSources:\n{source_list_str}", ""

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}")  # Log the history for debugging

    # Convert gradio history to the format expected by get_response_with_search
    chat_history = []
    if history:
        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})
            # Ignore any other formats

    try:
        full_response = ""
        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 only the new content
            new_content = main_content[len(full_response):]
            full_response = main_content
            yield new_content

        # Yield the sources as a separate message
        if sources:
            yield f"\n\nSources:\n{sources}"

    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)}"

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."

    async def wrapped_respond(*args):
        try:
            async for response in respond(*args):
                yield response
        except Exception as e:
            logging.error(f"Error in wrapped_respond: {str(e)}")
            yield f"An error occurred: {str(e)}"

    demo = gr.ChatInterface(
        fn=wrapped_respond,  # Use the wrapped version
        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,
        )
    )

    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)