skills/assistant/scripts/chat.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "pinecone==9.1.0",
#   "typer>=0.15.0",
#   "rich>=13.0.0",
# ]
# ///
"""
Chat with a Pinecone Assistant and receive cited responses.

Usage:
    uv run chat.py --assistant NAME --message "Your question" [--stream]

Environment Variables:
    PINECONE_API_KEY: Required Pinecone API key

Output:
    Assistant's response with citations to source documents
"""

import os
import typer
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from pinecone import Pinecone
from pinecone.models.assistant import Message

app = typer.Typer()
console = Console()


@app.command()
def main(
    assistant: str = typer.Option(..., "--assistant", "-a", help="Name of the assistant to chat with"),
    message: str = typer.Option(..., "--message", "-m", help="Your question or message"),
):
    """Chat with a Pinecone Assistant and receive answers with source citations."""

    # Check for API key
    api_key = os.environ.get("PINECONE_API_KEY")
    if not api_key:
        console.print("[red]Error: PINECONE_API_KEY environment variable not set[/red]")
        console.print("\nGet your API key from: https://app.pinecone.io/?sessionType=signup")
        raise typer.Exit(1)

    try:
        # Initialize Pinecone client
        pc = Pinecone(api_key=api_key,source_tag="codex_plugin:assistant")

        # Create message
        user_msg = Message(role="user", content=message)

        # Display user question
        console.print(Panel(f"[bold cyan]Question:[/bold cyan] {message}", border_style="cyan"))

        # Get response
        with console.status("[bold blue]Thinking...[/bold blue]"):
            response = pc.assistants.chat(assistant_name=assistant, messages=[user_msg], stream=False)

        answer_content = response.message.content
        citations = response.citations if hasattr(response, 'citations') else []
        usage = response.usage if hasattr(response, 'usage') else None

        # Display assistant's response (same for both modes)
        console.print("\n[bold green]Answer:[/bold green]\n")

        if answer_content:
            console.print(Panel(answer_content, border_style="green", title="Assistant Response"))
        else:
            console.print("[yellow]No response content received[/yellow]")

        # Display citations if available
        if citations and len(citations) > 0:

            console.print("\n[bold yellow]Citations:[/bold yellow]\n")

            citations_table = Table(show_header=True, header_style="bold yellow")
            citations_table.add_column("#", style="dim", width=4)
            citations_table.add_column("File", style="cyan", width=40)
            citations_table.add_column("Pages", style="blue", width=15)
            citations_table.add_column("Position", style="green", width=10)

            citation_num = 0
            for citation in citations:
                # Each citation has a list of references
                if hasattr(citation, 'references') and citation.references:
                    for reference in citation.references:
                        citation_num += 1

                        # Get file name
                        file_name = "Unknown"
                        if hasattr(reference, 'file') and hasattr(reference.file, 'name'):
                            file_name = reference.file.name

                        # Get pages
                        pages = []
                        if hasattr(reference, 'pages') and reference.pages:
                            pages = reference.pages

                        # Format pages
                        if pages:
                            pages_str = ", ".join(str(p) for p in pages)
                        else:
                            pages_str = "N/A"

                        # Get position from citation
                        position = getattr(citation, 'position', 'N/A')

                        citations_table.add_row(
                            str(citation_num),
                            file_name,
                            pages_str,
                            str(position)
                        )

            console.print(citations_table)

            # Optionally show download links
            console.print("\n[dim]Tip: File URLs are temporary signed links valid for ~1 hour[/dim]")

        # Display token usage
        if usage:
            usage_info = f"""[dim]Tokens used:[/dim]
• Prompt: {getattr(usage, 'prompt_tokens', 'N/A')}
• Completion: {getattr(usage, 'completion_tokens', 'N/A')}
• Total: {getattr(usage, 'total_tokens', 'N/A')}"""
            console.print(Panel(usage_info, border_style="dim", title="Usage Stats"))

        # Follow-up suggestion
        console.print(f"\n[dim]Continue the conversation with another message using the same command[/dim]")

    except Exception as e:
        console.print(f"[red]Error: {e}[/red]")
        raise typer.Exit(1)


if __name__ == "__main__":
    app()