Upload 3 files

pydantic_ai_agent.py

@@ -0,0 +1,142 @@
+from __future__ import annotations as _annotations
+import datetime
+from dataclasses import dataclass
+from dotenv import load_dotenv
+import logfire
+import asyncio
+import os
+from supabase import create_client
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.models.gemini import GeminiModel
+from supabase import Client
+from typing import List
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+
+load_dotenv()
+
+# Load API keys and initialize Supabase client
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+
+supabase: Client = create_client(
+    os.getenv("SUPABASE_URL1"),
+    os.getenv("SUPABASE_SERVICE_KEY1")
+)
+logfire.configure(send_to_logfire='always')
+
+@dataclass
+class PydanticAIDeps:
+    supabase: Client
+    gemini: GeminiModel
+
+# Updated system prompt explicitly instructs the agent to use the tools for documentation queries.
+system_prompt = """
+You are a PineScript expert with direct access to complete documentation through tool functions.
+Whenever a user asks for documentation, examples, or strategies for PineScript, you MUST:
+1. List all available documentation pages (using the tool "list_documentation_pages").
+2. Retrieve detailed, relevant documentation (using the tool "retrieve_relevant_documentation").
+Do not generate a direct answer from your internal data. 
+If no documentation is found, clearly state so.
+Then, combine the retrieved information to answer the query accurately.
+IMPORTANT: When a user asks "what is your name" or similar, always respond with "PineScript expert"
+"""
+
+model = GeminiModel('gemini-1.5-flash', api_key=GOOGLE_API_KEY)
+PineScript_expert = Agent(
+    model,
+    system_prompt=system_prompt,
+    deps_type=PydanticAIDeps,
+    retries=2
+)
+
+async def get_embedding(text: str) -> List[float]:
+    try:
+        embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", api_key=GOOGLE_API_KEY)
+        vector = embeddings.embed_query(text)
+        return vector
+    except Exception as e:
+        print(f"Error getting embedding: {e}")
+        return [0] * 768
+
+@PineScript_expert.tool
+async def retrieve_relevant_documentation(ctx: RunContext[PydanticAIDeps], user_query: str) -> str:
+    """
+    Retrieve relevant documentation chunks based on the query.
+    """
+    try:
+        print("DEBUG: retrieve_relevant_documentation tool function triggered.")
+        # Log the user query
+        log_result = ctx.deps.supabase.table('user_queries').insert({
+            'query': user_query,
+            'timestamp': datetime.datetime.now().isoformat()
+        }).execute()
+        error_info = getattr(log_result, "error", None)
+        if error_info:
+            print("Insert error:", error_info)
+        else:
+            print("Inserted user query successfully.")
+        # Get query embedding and retrieve matched documents
+        query_embedding = await get_embedding(user_query)
+        result = ctx.deps.supabase.rpc(
+            'match_site_pages',
+            {
+                'query_embedding': query_embedding,
+                'match_count': 20,
+                'filter': {'source': 'pydantic_ai_docs'}
+            }
+        ).execute()
+        if not result.data:
+            return "I couldn't find any relevant documentation for that query."
+        formatted_chunks = []
+        for doc in result.data:
+            chunks_text = f"""
+# {doc['title']}
+
+{doc['content']}
+"""
+            formatted_chunks.append(chunks_text)
+        return "\n\n---\n\n".join(formatted_chunks)
+    except Exception as e:
+        print(f"Error retrieving documentation: {e}")
+        return f"An error occurred while retrieving the documentation: {str(e)}"
+
+@PineScript_expert.tool
+async def list_documentation_pages(ctx: RunContext[PydanticAIDeps]) -> List[str]:
+    """
+    List all available PineScript documentation pages.
+    """
+    try:
+        result = ctx.deps.supabase.from_('site_pages') \
+            .select('url') \
+            .eq('metadata->>source', 'pydantic_ai_docs') \
+            .execute()
+        if not result.data:
+            return []
+        urls = sorted(set(doc['url'] for doc in result.data))
+        return urls
+    except Exception as e:
+        print(f"Error retrieving documentation pages: {e}")
+        return []
+
+@PineScript_expert.tool
+async def get_page_content(ctx: RunContext[PydanticAIDeps], url: str) -> str:
+    """
+    Retrieve the full content of a specific documentation page.
+    """
+    try:
+        result = ctx.deps.supabase.from_('site_pages') \
+            .select('title, content, chunk_number') \
+            .eq('url', url) \
+            .eq('metadata->>source', 'pydantic_ai_docs') \
+            .order('chunk_number') \
+            .execute()
+        if not result.data:
+            return f"No content found for URL: {url}"
+        page_title = result.data[0]['title'].split(' - ')[0]
+        formatted_content = [f"# {page_title}\n"]
+        for chunk in result.data:
+            formatted_content.append(chunk['content'])
+        return "\n\n".join(formatted_content)
+    except Exception as e:
+        print(f"Error retrieving page content: {e}")
+        return f"Error retrieving page content: {str(e)}"

streamlit_ui.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+from typing import Literal, TypedDict
+import asyncio
+import os
+import streamlit as st
+import json
+import datetime
+import logfire
+from supabase import Client
+from pydantic_ai.models.gemini import GeminiModel
+from supabase import create_client
+
+# Import all the message part classes
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+    UserPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    RetryPromptPart,
+    ModelMessagesTypeAdapter
+)
+from pydantic_ai_agent import PineScript_expert, PydanticAIDeps
+
+# Load environment variables
+from dotenv import load_dotenv
+load_dotenv()
+
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+
+# Initialize Supabase client
+supabase: Client = create_client(
+    os.getenv("SUPABASE_URL1"),
+    os.getenv("SUPABASE_SERVICE_KEY1")
+)
+
+# Configure logfire to suppress warnings (optional)
+logfire.configure(send_to_logfire='always')
+
+class ChatMessage(TypedDict):
+    """Format of messages sent to the browser/API."""
+    role: Literal['user', 'model']
+    timestamp: str
+    content: str
+
+def display_message_part(part):
+    """
+    Display a single part of a message in the Streamlit UI.
+    Customize how you display system prompts, user prompts,
+    tool calls, tool returns, etc.
+    """
+    # system-prompt
+    if part.part_kind == 'system-prompt':
+        with st.chat_message("system"):
+            st.markdown(f"**System**: {part.content}")
+    # user-prompt
+    elif part.part_kind == 'user-prompt':
+        with st.chat_message("user"):
+            st.markdown(part.content)
+    # text
+    elif part.part_kind == 'text':
+        with st.chat_message("assistant"):
+            st.markdown(part.content)
+
+async def run_agent_with_streaming(user_input: str):
+    """
+    Run the agent with streaming text for the user_input prompt,
+    while maintaining the entire conversation in `st.session_state.messages`.
+    """
+    # Prepare dependencies for the agent
+    deps = PydanticAIDeps(
+        supabase=supabase,
+        gemini=GeminiModel('gemini-1.5-flash', api_key=GOOGLE_API_KEY)
+    )
+    message_placeholder = st.empty()
+    partial_text = ""
+    try:
+        # Run the agent in a stream
+        async with PineScript_expert.run_stream(
+            user_input,
+            deps=deps,
+            message_history=st.session_state.messages[:-1],  # pass entire conversation so far
+        ) as result:
+            # Render partial text as it arrives
+            async for chunk in result.stream_text(delta=True):
+                partial_text += chunk
+                message_placeholder.markdown(partial_text)
+    except Exception as e:
+        # Instead of displaying the raw error, display a friendly message.
+        message_placeholder.markdown("**Could you please ask again?**")
+        st.error("Could you please ask again?")
+        # Also, add the friendly message to the conversation history.
+        st.session_state.messages.append(
+            ModelResponse(parts=[TextPart(content="Could you please ask again?")])
+        )
+        return
+
+    # Now that the stream is finished, add new messages from this run
+    filtered_messages = [
+        msg for msg in result.new_messages()
+        if not (hasattr(msg, 'parts') and any(part.part_kind == 'user-prompt' for part in msg.parts))
+    ]
+    st.session_state.messages.extend(filtered_messages)
+    st.session_state.messages.append(
+        ModelResponse(parts=[TextPart(content=partial_text)])
+    )
+
+
+def store_user_query(user_query: str):
+    """
+    Stores the user's query in the 'user_queries' table.
+    """
+    try:
+        result = supabase.table('user_queries').insert({
+            'query': user_query,
+            'timestamp': datetime.datetime.now().isoformat()
+        }).execute()
+        print("User query stored successfully.")
+    except Exception as e:
+        print(f"Error storing user query: {e}")
+
+async def main():
+    st.title("PineScript Agentic RAG")
+    st.write("Ask any question about PineScript, the hidden truths of its beauty lie within.")
+
+    # Initialize chat history in session state if not present
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    # Display all messages from the conversation so far
+    # Each message is either a ModelRequest or ModelResponse.
+    # We iterate over their parts to decide how to display them.
+    for msg in st.session_state.messages:
+        if isinstance(msg, ModelRequest) or isinstance(msg, ModelResponse):
+            for part in msg.parts:
+                display_message_part(part)
+
+    # Chat input for the user
+    user_input = st.chat_input("What questions do you have about PineScript?")
+
+    if user_input:
+        # We append a new request to the conversation explicitly
+        store_user_query(user_input)
+        st.session_state.messages.append(
+            ModelRequest(parts=[UserPromptPart(content=user_input)])
+        )
+        
+        # Display user prompt in the UI
+        with st.chat_message("user"):
+            st.markdown(user_input)
+
+        # Display the assistant's partial response while streaming
+        with st.chat_message("assistant"):
+            # Actually run the agent now, streaming the text
+            await run_agent_with_streaming(user_input)
+
+if __name__ == "__main__":
+    asyncio.run(main())