add code

.gitignore

@@ -0,0 +1,5 @@
+.env
+db_creation.py
+db_test.py
+docs.py
+__pycache__/

app.py

@@ -0,0 +1,212 @@
+from __future__ import annotations as _annotations
+
+import json
+import os
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import AsyncGenerator
+
+import asyncpg
+import gradio as gr
+import numpy as np
+import pydantic_core
+from gradio_webrtc import (
+    AdditionalOutputs,
+    ReplyOnPause,
+    WebRTC,
+    audio_to_bytes,
+    get_twilio_turn_credentials,
+)
+from groq import Groq
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+from pydantic_ai import RunContext
+from pydantic_ai.agent import Agent
+from pydantic_ai.messages import ModelStructuredResponse, ModelTextResponse, ToolReturn
+
+DOCS = json.load(open("gradio_docs.json"))
+
+groq_client = Groq()
+openai = AsyncOpenAI()
+
+
+@dataclass
+class Deps:
+    openai: AsyncOpenAI
+    pool: asyncpg.Pool
+
+
+SYSTEM_PROMPT = (
+    "You are an assistant designed to help users answer questions about Gradio. "
+    "You have a retrival tool that can provide relevant documentation sections based on the user query. "
+    "Be curteous and helpful to the user but feel free to refuse answering questions that are not about Gradio. "
+)
+
+
+agent = Agent(
+    "openai:gpt-4o",
+    deps_type=Deps,
+    system_prompt=SYSTEM_PROMPT,
+)
+
+
+class RetrievalResult(BaseModel):
+    content: str
+    ids: list[int]
+
+
+@asynccontextmanager
+async def database_connect(
+    create_db: bool = False,
+) -> AsyncGenerator[asyncpg.Pool, None]:
+    server_dsn, database = (
+        os.getenv("DATABASE_URL"),
+        "gradio_ai_rag",
+    )
+    if create_db:
+        conn = await asyncpg.connect(server_dsn)
+        try:
+            db_exists = await conn.fetchval(
+                "SELECT 1 FROM pg_database WHERE datname = $1", database
+            )
+            if not db_exists:
+                await conn.execute(f"CREATE DATABASE {database}")
+        finally:
+            await conn.close()
+
+    pool = await asyncpg.create_pool(f"{server_dsn}/{database}")
+    try:
+        yield pool
+    finally:
+        await pool.close()
+
+
+@agent.tool
+async def retrieve(context: RunContext[Deps], search_query: str) -> str:
+    """Retrieve documentation sections based on a search query.
+
+    Args:
+        context: The call context.
+        search_query: The search query.
+    """
+    print(f"create embedding for {search_query}")
+    embedding = await context.deps.openai.embeddings.create(
+        input=search_query,
+        model="text-embedding-3-small",
+    )
+
+    assert (
+        len(embedding.data) == 1
+    ), f"Expected 1 embedding, got {len(embedding.data)}, doc query: {search_query!r}"
+    embedding = embedding.data[0].embedding
+    embedding_json = pydantic_core.to_json(embedding).decode()
+    rows = await context.deps.pool.fetch(
+        "SELECT id, title, content FROM doc_sections ORDER BY embedding <-> $1 LIMIT 8",
+        embedding_json,
+    )
+    content = "\n\n".join(f'# {row["title"]}\n{row["content"]}\n' for row in rows)
+    ids = [row["id"] for row in rows]
+    return RetrievalResult(content=content, ids=ids).model_dump_json()
+
+
+async def stream_from_agent(
+    audio: tuple[int, np.ndarray], chatbot: list[dict], past_messages: list
+):
+    question = groq_client.audio.transcriptions.create(
+        file=("audio-file.mp3", audio_to_bytes(audio)),
+        model="whisper-large-v3-turbo",
+        response_format="verbose_json",
+    ).text
+
+    print("text", question)
+
+    chatbot.append({"role": "user", "content": question})
+    yield AdditionalOutputs(chatbot, gr.skip())
+
+    async with database_connect(False) as pool:
+        deps = Deps(openai=openai, pool=pool)
+        async with agent.run_stream(
+            question, deps=deps, message_history=past_messages
+        ) as result:
+            for message in result.new_messages():
+                past_messages.append(message)
+                if isinstance(message, ModelStructuredResponse):
+                    for call in message.calls:
+                        gr_message = {
+                            "role": "assistant",
+                            "content": "",
+                            "metadata": {
+                                "title": "🔍 Retrieving relevant docs",
+                                "id": call.tool_id,
+                            },
+                        }
+                        chatbot.append(gr_message)
+                if isinstance(message, ToolReturn):
+                    for gr_message in chatbot:
+                        if (
+                            gr_message.get("metadata", {}).get("id", "")
+                            == message.tool_id
+                        ):
+                            paths = []
+                            for d in DOCS:
+                                tool_result = RetrievalResult.model_validate_json(
+                                    message.content
+                                )
+                                if d["id"] in tool_result.ids:
+                                    paths.append(d["path"])
+                            gr_message["content"] = (
+                                f"Relevant Context:\n {'\n'.join(list(set(paths)))}"
+                            )
+                yield AdditionalOutputs(chatbot, gr.skip())
+            chatbot.append({"role": "assistant", "content": ""})
+            async for message in result.stream_text():
+                chatbot[-1]["content"] = message
+                yield AdditionalOutputs(chatbot, gr.skip())
+            data = await result.get_data()
+            past_messages.append(ModelTextResponse(content=data))
+            yield AdditionalOutputs(gr.skip(), past_messages)
+
+
+with gr.Blocks() as demo:
+    placeholder = """
+<div style="display: flex; justify-content: center; align-items: center; gap: 1rem; padding: 1rem; width: 100%">
+    <img src="/gradio_api/file=logo.svg" style="max-width: 200px; height: auto">
+    <div>
+        <h1 style="margin: 0 0 1rem 0">Chat with Gradio Docs 🗣️</h1>
+        <h3 style="margin: 0 0 0.5rem 0">
+            Simple RAG agent over Gradio docs built with Pydantic AI.
+        </h3>
+        <h3 style="margin: 0">
+            Ask any question about Gradio with your natural voice and get an answer!
+        </h3>
+    </div>
+</div>
+"""
+    past_messages = gr.State([])
+    chatbot = gr.Chatbot(
+        label="Gradio Docs Bot",
+        type="messages",
+        placeholder=placeholder,
+        avatar_images=(None, "logo.svg"),
+    )
+    audio = WebRTC(
+        label="Talk with the Agent",
+        modality="audio",
+        rtc_configuration=get_twilio_turn_credentials(),
+        mode="send",
+    )
+    audio.stream(
+        ReplyOnPause(stream_from_agent),
+        inputs=[audio, chatbot, past_messages],
+        outputs=[audio],
+    )
+    audio.on_additional_outputs(
+        lambda c, s: (c, s),
+        outputs=[chatbot, past_messages],
+        queue=False,
+        show_progress="hidden",
+    )
+
+
+if __name__ == "__main__":
+    demo.launch(allowed_paths=["logo.svg"])

requirements.in

@@ -0,0 +1,6 @@
+gradio-webrtc[vad]>=0.0.18
+numba>=0.60.0
+pydantic-ai
+asyncpg
+groq
+openai

requirements.txt

@@ -0,0 +1,296 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile requirements.in -o requirements.txt
+aiofiles==23.2.1
+    # via gradio
+aioice==0.9.0
+    # via aiortc
+aiortc==1.9.0
+    # via gradio-webrtc
+annotated-types==0.7.0
+    # via pydantic
+anyio==4.7.0
+    # via
+    #   gradio
+    #   groq
+    #   httpx
+    #   openai
+    #   starlette
+asyncpg==0.30.0
+    # via -r requirements.in
+audioread==3.0.1
+    # via librosa
+av==12.3.0
+    # via aiortc
+cachetools==5.5.0
+    # via google-auth
+certifi==2024.8.30
+    # via
+    #   httpcore
+    #   httpx
+    #   requests
+cffi==1.17.1
+    # via
+    #   aiortc
+    #   cryptography
+    #   pylibsrtp
+    #   soundfile
+charset-normalizer==3.4.0
+    # via requests
+click==8.1.7
+    # via
+    #   typer
+    #   uvicorn
+colorama==0.4.6
+    # via griffe
+coloredlogs==15.0.1
+    # via onnxruntime
+cryptography==44.0.0
+    # via
+    #   aiortc
+    #   pyopenssl
+decorator==5.1.1
+    # via librosa
+distro==1.9.0
+    # via
+    #   groq
+    #   openai
+dnspython==2.7.0
+    # via aioice
+eval-type-backport==0.2.0
+    # via pydantic-ai-slim
+fastapi==0.115.6
+    # via gradio
+ffmpy==0.4.0
+    # via gradio
+filelock==3.16.1
+    # via huggingface-hub
+flatbuffers==24.3.25
+    # via onnxruntime
+fsspec==2024.10.0
+    # via
+    #   gradio-client
+    #   huggingface-hub
+google-auth==2.36.0
+    # via pydantic-ai-slim
+google-crc32c==1.6.0
+    # via aiortc
+gradio==5.8.0
+    # via gradio-webrtc
+gradio-client==1.5.1
+    # via gradio
+gradio-webrtc==0.0.18
+    # via -r requirements.in
+griffe==1.5.1
+    # via pydantic-ai-slim
+groq==0.13.0
+    # via
+    #   -r requirements.in
+    #   pydantic-ai-slim
+h11==0.14.0
+    # via
+    #   httpcore
+    #   uvicorn
+httpcore==1.0.7
+    # via httpx
+httpx==0.28.0
+    # via
+    #   gradio
+    #   gradio-client
+    #   groq
+    #   openai
+    #   pydantic-ai-slim
+    #   safehttpx
+huggingface-hub==0.26.3
+    # via
+    #   gradio
+    #   gradio-client
+humanfriendly==10.0
+    # via coloredlogs
+idna==3.10
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+ifaddr==0.2.0
+    # via aioice
+jinja2==3.1.4
+    # via gradio
+jiter==0.8.0
+    # via openai
+joblib==1.4.2
+    # via
+    #   librosa
+    #   scikit-learn
+lazy-loader==0.4
+    # via librosa
+librosa==0.10.2.post1
+    # via gradio-webrtc
+llvmlite==0.43.0
+    # via numba
+logfire-api==2.6.2
+    # via pydantic-ai-slim
+markdown-it-py==3.0.0
+    # via rich
+markupsafe==2.1.5
+    # via
+    #   gradio
+    #   jinja2
+mdurl==0.1.2
+    # via markdown-it-py
+mpmath==1.3.0
+    # via sympy
+msgpack==1.1.0
+    # via librosa
+numba==0.60.0
+    # via
+    #   -r requirements.in
+    #   librosa
+numpy==2.0.2
+    # via
+    #   gradio
+    #   librosa
+    #   numba
+    #   onnxruntime
+    #   pandas
+    #   scikit-learn
+    #   scipy
+    #   soxr
+onnxruntime==1.20.1
+    # via gradio-webrtc
+openai==1.57.0
+    # via
+    #   -r requirements.in
+    #   pydantic-ai-slim
+orjson==3.10.12
+    # via gradio
+packaging==24.2
+    # via
+    #   gradio
+    #   gradio-client
+    #   huggingface-hub
+    #   lazy-loader
+    #   onnxruntime
+    #   pooch
+pandas==2.2.3
+    # via gradio
+pillow==11.0.0
+    # via gradio
+platformdirs==4.3.6
+    # via pooch
+pooch==1.8.2
+    # via librosa
+protobuf==5.29.1
+    # via onnxruntime
+pyasn1==0.6.1
+    # via
+    #   pyasn1-modules
+    #   rsa
+pyasn1-modules==0.4.1
+    # via google-auth
+pycparser==2.22
+    # via cffi
+pydantic==2.10.3
+    # via
+    #   fastapi
+    #   gradio
+    #   groq
+    #   openai
+    #   pydantic-ai-slim
+pydantic-ai==0.0.9
+    # via -r requirements.in
+pydantic-ai-slim==0.0.9
+    # via pydantic-ai
+pydantic-core==2.27.1
+    # via pydantic
+pydub==0.25.1
+    # via gradio
+pyee==12.1.1
+    # via aiortc
+pygments==2.18.0
+    # via rich
+pylibsrtp==0.10.0
+    # via aiortc
+pyopenssl==24.3.0
+    # via aiortc
+python-dateutil==2.9.0.post0
+    # via pandas
+python-multipart==0.0.19
+    # via gradio
+pytz==2024.2
+    # via pandas
+pyyaml==6.0.2
+    # via
+    #   gradio
+    #   huggingface-hub
+requests==2.32.3
+    # via
+    #   huggingface-hub
+    #   pooch
+    #   pydantic-ai-slim
+rich==13.9.4
+    # via typer
+rsa==4.9
+    # via google-auth
+ruff==0.8.2
+    # via gradio
+safehttpx==0.1.6
+    # via gradio
+scikit-learn==1.5.2
+    # via librosa
+scipy==1.14.1
+    # via
+    #   librosa
+    #   scikit-learn
+semantic-version==2.10.0
+    # via gradio
+shellingham==1.5.4
+    # via typer
+six==1.17.0
+    # via python-dateutil
+sniffio==1.3.1
+    # via
+    #   anyio
+    #   groq
+    #   openai
+soundfile==0.12.1
+    # via librosa
+soxr==0.5.0.post1
+    # via librosa
+starlette==0.41.3
+    # via
+    #   fastapi
+    #   gradio
+sympy==1.13.3
+    # via onnxruntime
+threadpoolctl==3.5.0
+    # via scikit-learn
+tomlkit==0.13.2
+    # via gradio
+tqdm==4.67.1
+    # via
+    #   huggingface-hub
+    #   openai
+typer==0.15.1
+    # via gradio
+typing-extensions==4.12.2
+    # via
+    #   anyio
+    #   fastapi
+    #   gradio
+    #   gradio-client
+    #   groq
+    #   huggingface-hub
+    #   librosa
+    #   openai
+    #   pydantic
+    #   pydantic-core
+    #   pyee
+    #   typer
+tzdata==2024.2
+    # via pandas
+urllib3==2.2.3
+    # via requests
+uvicorn==0.32.1
+    # via gradio
+websockets==14.1
+    # via gradio-client