Setting the file. One moment. (opens in a new tab)
scripts/context.py
Python·151 lines·6 KB
PINECONE_API_KEY: Required Pinecone API key
18
19Output:
20 Relevant context snippets with file sources, page numbers, and relevance scores
21"""
22
23import os
24import json as json_module
25import typer
26from rich.console import Console
27from rich.panel import Panel
28from rich.table import Table
29from rich.text import Text
30from pinecone import Pinecone
31from pinecone.models.assistant import TextSnippet
32
33app = typer.Typer()
34console = Console()
35
36
37@app.command()
38def main(
39 assistant: str = typer.Option(..., "--assistant", "-a", help="Name of the assistant"),
40 query: str = typer.Option(..., "--query", "-q", help="Search query text"),
41 top_k: int = typer.Option(5, "--top-k", "-k", help="Number of results to return (max 16)"),
42 snippet_size: int = typer.Option(1024, "--snippet-size", "-s", help="Maximum tokens per snippet"),
43 json: bool = typer.Option(False, "--json", help="Output in JSON format"),
44):
45 """Retrieve relevant context snippets from an assistant's knowledge base."""
46
47 # Check for API key
48 api_key = os.environ.get("PINECONE_API_KEY")
49 if not api_key:
50 console.print("[red]Error: PINECONE_API_KEY environment variable not set[/red]")
51 console.print("\nGet your API key from: https://app.pinecone.io/?sessionType=signup")
52 raise typer.Exit(1)
53
54 try:
55 # Initialize Pinecone client
56 pc = Pinecone(api_key=api_key, source_tag="pinecone_skills:assistant")
57
58 # Display query
59 if not json:
60 console.print(Panel(f"[bold cyan]Query:[/bold cyan] {query}", border_style="cyan"))
61
62 # Retrieve context
63 with console.status("[bold blue]Searching knowledge base...[/bold blue]", spinner="dots"):
64 response = pc.assistants.context(
65 assistant_name=assistant, query=query, top_k=top_k, snippet_size=snippet_size
66 )
67
68 # Get snippets from response
69 snippets = response.snippets if hasattr(response, 'snippets') else []
70
71 if json:
72 # JSON output
73 results = []
74 for snippet in snippets:
75 file_name = "Unknown"
76 pages = []
77 if hasattr(snippet, 'reference') and snippet.reference:
78 ref = snippet.reference
79 if hasattr(ref, 'file') and hasattr(ref.file, 'name'):
80 file_name = ref.file.name
81 if hasattr(ref, 'pages') and ref.pages:
82 pages = ref.pages
83
84 results.append({
85 "file_name": file_name,
86 "pages": pages,
87 "content": getattr(snippet, 'content', ''),
88 "score": getattr(snippet, 'score', 0.0),
89 # SDK 9 encodes the snippet kind as a msgspec tag, not an
90 # instance attribute, so getattr(snippet, 'type') silently
91 # returned the default for every snippet.
92 "type": "text" if isinstance(snippet, TextSnippet) else "multimodal",
93 })
94 print(json_module.dumps({"snippets": results, "count": len(results)}, indent=2))
95 else:
96 # Rich formatted output
97 if not snippets or len(snippets) == 0:
98 console.print("[yellow]No context found for this query[/yellow]")
99 return
100
101 console.print(f"\n[bold]Found {len(snippets)} relevant snippet(s):[/bold]\n")
102
103 for idx, snippet in enumerate(snippets, 1):
104 # Extract file info from reference
105 file_name = "Unknown"
106 pages = []
107 if hasattr(snippet, 'reference') and snippet.reference:
108 ref = snippet.reference
109 if hasattr(ref, 'file') and hasattr(ref.file, 'name'):
110 file_name = ref.file.name
111 if hasattr(ref, 'pages') and ref.pages:
112 pages = ref.pages
113
114 score = getattr(snippet, 'score', 0.0)
115 content = getattr(snippet, 'content', '')
116
117 # Create header
118 header = f"#{idx} - {file_name}"
119 if pages:
120 pages_str = ", ".join(str(p) for p in pages)
121 header += f" (Page {pages_str})"
122 header += f" - Score: {score:.3f}" if isinstance(score, (int, float)) else f" - Score: {score}"
123
124 console.print(Panel(
125 content,
126 title=header,
127 border_style="blue",
128 subtitle=f"[dim]Relevance: {score:.2%}[/dim]" if isinstance(score, (int, float)) else None
129 ))
130 console.print()
131
132 # Suggest next action
133 next_action = f"""[bold]Next steps:[/bold]
134<<next_after_context>>"""
135 console.print(Panel(next_action, title="What's Next?", border_style="green"))
136
137 except AttributeError as e:
138 # Handle case where context method doesn't exist or response structure is different
139 console.print(f"[red]Error: Context retrieval failed[/red]")
140 console.print(f"[dim]Details: {e}[/dim]")
141 console.print("\n[yellow]Note:[/yellow] Context API requires SDK version with assistant.context() support")
142 console.print("\n[yellow]Try using chat instead:[/yellow]")
143 console.print(f" <<next_chat_fallback>>")
144 raise typer.Exit(1)
145 except Exception as e:
146 console.print(f"[red]Error: {e}[/red]")
147 raise typer.Exit(1)
148
149
150if __name__ == "__main__":
151 app()