app.py

import os
import time
import gradio as gr

from groq import Groq
from pinecone import Pinecone, ServerlessSpec
from semantic_router.encoders import HuggingFaceEncoder

encoder = HuggingFaceEncoder(name="dwzhu/e5-base-4k")
embeds = encoder(["this is a test"])
dims = len(embeds[0])

groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
index_name = "groq-llama-3-rag"

if index_name not in [i["name"] for i in pc.list_indexes()]:
    pc.create_index(index_name, dimension=dims, metric='cosine', spec=ServerlessSpec(cloud="aws", region="us-east-1"))
    while not pc.describe_index(index_name).status['ready']:
        time.sleep(1)

index = pc.Index(index_name)

def get_docs(query: str, top_k: int) -> list[str]:
    xq = encoder([query])
    res = index.query(vector=xq, top_k=top_k, include_metadata=True)
    return [x["metadata"]['content_snippet'] for x in res["matches"]]

def generate(query: str, history):
    if not history:
        system_message = (
            "You are a friendly and knowledgeable New Yorker who loves sharing recommendations about the city. "
            "You have lived in NYC for years and know both the famous tourist spots and hidden local gems. "
            "Your goal is to give recommendations tailored to what the user is asking for, whether they want iconic attractions "
            "or lesser-known spots loved by locals.\n\n"
            "Use the provided context to enhance your responses with real local insights, but only include details that are relevant "
            "to the user’s question.\n\n"
            "CONTEXT:\n" + "\n---\n".join(get_docs(query, top_k=5))
        )
        messages = [{"role": "system", "content": system_message}]
    else:
        messages = []
        for user_msg, bot_msg in history:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": bot_msg})
        messages.append({"role": "system", "content": "CONTEXT:\n" + "\n---\n".join(get_docs(query, top_k=5))})

    messages.append({"role": "user", "content": query})
    response = groq_client.chat.completions.create(model="llama3-70b-8192", messages=messages)
    return response.choices[0].message.content

custom_css = """
.gradio-container {
    background: #ffffff !important;
}
.chatbot {
    background: #f8f8f8 !important;
    border-radius: 15px;
    padding: 10px;
}
.chat-message {
    display: flex;
    align-items: center;
    margin-bottom: 10px;
}
.chat-message.user {
    justify-content: flex-end;
}
.chat-message.assistant {
    justify-content: flex-start;
}
.chat-bubble {
    padding: 10px 15px;
    border-radius: 20px;
    max-width: 70%;
    font-size: 16px;
    display: inline-block;
}
.chat-bubble.user {
    background-color: #007aff;
    color: white;
    border-bottom-right-radius: 5px;
}
.chat-bubble.assistant {
    background-color: #d1d1d6;
    color: black;
    border-bottom-left-radius: 5px;
}
"""

with gr.Blocks(css=custom_css) as demo:
    chatbot = gr.Chatbot(label="NYC Buddy", elem_classes=["chatbot"])
    state = gr.State([])
    user_input = gr.Textbox(placeholder="Ask me anything about NYC!")
    send_btn = gr.Button("Send")

    def respond(message, history):
        response = generate(message, history)
        history = history + [(message, response)]
        return history, ""

    def startup():
        welcome_msg = "Hey! I'm your NYC Buddy. Looking for local tips, hidden gems, or iconic spots? Just ask!"
        return [(None, welcome_msg)], []

    demo.load(startup, outputs=[chatbot, state])
    send_btn.click(respond, inputs=[user_input, state], outputs=[chatbot, user_input])

demo.launch()