Setting the file. One moment. (opens in a new tab)
scripts/chat.py
Python·141 lines·5 KB
PINECONE_API_KEY: Required Pinecone API key
18
19Output:
20 Assistant's response with citations to source documents
21"""
22
23import os
24import typer
25from rich.console import Console
26from rich.panel import Panel
27from rich.table import Table
28from pinecone import Pinecone
29from pinecone.models.assistant import Message
30
31app = typer.Typer()
32console = Console()
33
34
35@app.command()
36def main(
37 assistant: str = typer.Option(..., "--assistant", "-a", help="Name of the assistant to chat with"),
38 message: str = typer.Option(..., "--message", "-m", help="Your question or message"),
39):
40 """Chat with a Pinecone Assistant and receive answers with source citations."""
41
42 # Check for API key
43 api_key = os.environ.get("PINECONE_API_KEY")
44 if not api_key:
45 console.print("[red]Error: PINECONE_API_KEY environment variable not set[/red]")
46 console.print("\nGet your API key from: https://app.pinecone.io/?sessionType=signup")
47 raise typer.Exit(1)
48
49 try:
50 # Initialize Pinecone client
51 pc = Pinecone(api_key=api_key,source_tag="pinecone_skills:assistant")
52
53 # Create message
54 user_msg = Message(role="user", content=message)
55
56 # Display user question
57 console.print(Panel(f"[bold cyan]Question:[/bold cyan] {message}", border_style="cyan"))
58
59 # Get response
60 with console.status("[bold blue]Thinking...[/bold blue]"):
61 response = pc.assistants.chat(assistant_name=assistant, messages=[user_msg], stream=False)
62
63 answer_content = response.message.content
64 citations = response.citations if hasattr(response, 'citations') else []
65 usage = response.usage if hasattr(response, 'usage') else None
66
67 # Display assistant's response (same for both modes)
68 console.print("\n[bold green]Answer:[/bold green]\n")
69
70 if answer_content:
71 console.print(Panel(answer_content, border_style="green", title="Assistant Response"))
72 else:
73 console.print("[yellow]No response content received[/yellow]")
74
75 # Display citations if available
76 if citations and len(citations) > 0:
77
78 console.print("\n[bold yellow]Citations:[/bold yellow]\n")
79
80 citations_table = Table(show_header=True, header_style="bold yellow")
81 citations_table.add_column("#", style="dim", width=4)
82 citations_table.add_column("File", style="cyan", width=40)
83 citations_table.add_column("Pages", style="blue", width=15)
84 citations_table.add_column("Position", style="green", width=10)
85
86 citation_num = 0
87 for citation in citations:
88 # Each citation has a list of references
89 if hasattr(citation, 'references') and citation.references:
90 for reference in citation.references:
91 citation_num += 1
92
93 # Get file name
94 file_name = "Unknown"
95 if hasattr(reference, 'file') and hasattr(reference.file, 'name'):
96 file_name = reference.file.name
97
98 # Get pages
99 pages = []
100 if hasattr(reference, 'pages') and reference.pages:
101 pages = reference.pages
102
103 # Format pages
104 if pages:
105 pages_str = ", ".join(str(p) for p in pages)
106 else:
107 pages_str = "N/A"
108
109 # Get position from citation
110 position = getattr(citation, 'position', 'N/A')
111
112 citations_table.add_row(
113 str(citation_num),
114 file_name,
115 pages_str,
116 str(position)
117 )
118
119 console.print(citations_table)
120
121 # Optionally show download links
122 console.print("\n[dim]Tip: File URLs are temporary signed links valid for ~1 hour[/dim]")
123
124 # Display token usage
125 if usage:
126 usage_info = f"""[dim]Tokens used:[/dim]
127• Prompt: {getattr(usage, 'prompt_tokens', 'N/A')}
128• Completion: {getattr(usage, 'completion_tokens', 'N/A')}
129• Total: {getattr(usage, 'total_tokens', 'N/A')}"""
130 console.print(Panel(usage_info, border_style="dim", title="Usage Stats"))
131
132 # Follow-up suggestion
133 console.print(f"\n[dim]Continue the conversation with another message using the same command[/dim]")
134
135 except Exception as e:
136 console.print(f"[red]Error: {e}[/red]")
137 raise typer.Exit(1)
138
139
140if __name__ == "__main__":
141 app()