Skip to content

Commit aef1ba2

Browse files
muldercwmuldercw
committed
feat: Enhance chat assistant with full codebase indexing and graceful error handling
- Index entire clarifai package (not just cli/client/utils) - Add graceful error handling for model unavailability - Improve RAG system to handle all CLI-related questions - Support conversation history queries - Add rich markdown rendering for responses - Support chat history commands (history, clear) - Better system prompts for pipeline/model/deployment CLI support - Add rich library dependency for formatted output
1 parent 9b5d67a commit aef1ba2

3 files changed

Lines changed: 143 additions & 64 deletions

File tree

‎clarifai/cli/chat.py‎

Lines changed: 88 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
import sys
55

66
import click
7+
from rich.console import Console
8+
from rich.markdown import Markdown
79

810
import clarifai
911
from clarifai.cli.base import cli
@@ -15,11 +17,14 @@
1517
# Default model URL for GPT-OSS-120B chat completion
1618
DEFAULT_CHAT_MODEL_URL = "https://clarifai.com/openai/chat-completion/models/gpt-oss-120b"
1719

20+
# Rich console for formatted output
21+
console = Console()
22+
1823

1924
@cli.command()
2025
@click.pass_context
2126
def chat(ctx):
22-
"""Start an interactive chat session using the Clarifai chat model.
27+
"""Start an interactive session using the Clarifai CLI assistant.
2328
2429
You must be logged in first. Uses the current context's credentials.
2530
@@ -42,8 +47,9 @@ def chat(ctx):
4247
chat_model_url = current_context.get('chat_model_url', DEFAULT_CHAT_MODEL_URL)
4348

4449
# Try to initialize the model
50+
model_available = True
4551
try:
46-
click.secho("Initializing chat model...", fg='cyan')
52+
click.secho("Initializing assistant...", fg='cyan')
4753

4854
# Initialize model with the current context's PAT
4955
model = Model(
@@ -58,9 +64,9 @@ def chat(ctx):
5864
# Find the clarifai-python root
5965
clarifai_root = os.path.dirname(os.path.dirname(clarifai.__file__))
6066
rag = ClarifaiCodeRAG(clarifai_root)
61-
click.secho("(CLI reference system ready)\n", fg='blue')
67+
click.secho("(Knowledge base ready.)\n", fg='blue')
6268
except Exception as e:
63-
click.secho(f"(Warning: Could not load CLI reference system: {e})\n", fg='yellow')
69+
click.secho(f"(Warning: Could not load knowledge base: {e})\n", fg='yellow')
6470
rag = None
6571

6672
# Interactive chat loop
@@ -76,36 +82,79 @@ def chat(ctx):
7682
if not user_input:
7783
continue
7884

79-
# Check for exit commands
85+
# Check for special commands
8086
if user_input.lower() in ('exit', 'quit', 'bye'):
8187
click.secho("Goodbye!", fg='yellow')
8288
break
8389

90+
if user_input.lower() in ('help', '?'):
91+
click.echo(
92+
"\nChat Commands:\n"
93+
" exit, quit, bye - Exit the chat\n"
94+
" history, hist - Show conversation history\n"
95+
" clear - Clear conversation history\n"
96+
" help, ? - Show this help message\n"
97+
)
98+
continue
99+
100+
if user_input.lower() in ('history', 'hist'):
101+
if not conversation_history:
102+
click.echo("No conversation history yet.\n")
103+
else:
104+
click.echo("\nConversation History:")
105+
for i, msg in enumerate(conversation_history, 1):
106+
role = "You" if msg['role'] == 'user' else "🤖 Assistant"
107+
click.echo(f"\n{i}. {role}:")
108+
click.echo(
109+
f" {msg['message'][:300]}..."
110+
if len(msg['message']) > 300
111+
else f" {msg['message']}"
112+
)
113+
click.echo()
114+
continue
115+
116+
if user_input.lower() == 'clear':
117+
conversation_history = []
118+
click.secho("Conversation history cleared.\n", fg='green')
119+
continue
120+
84121
# Add to conversation history
85122
conversation_history.append({'role': 'user', 'message': user_input})
86123

87-
click.secho("Thinking...", fg='yellow')
124+
# click.secho("Thinking...", fg='yellow')
88125

89126
# Build conversation context for follow-up questions
90127
conversation_context = ""
91-
if len(conversation_history) > 2: # Include last exchange if available
92-
conversation_context = "\n## Previous Context:\n"
93-
# Include last 2 exchanges (4 messages max) for context
94-
for msg in conversation_history[-4:-1]:
128+
if len(conversation_history) > 1:
129+
conversation_context = "\n## Conversation History:\n"
130+
# Include all previous messages (not just last 4) for full context
131+
for msg in conversation_history[:-1]: # Exclude the current user message
95132
role = "User" if msg['role'] == 'user' else "Assistant"
96-
conversation_context += (
97-
f"{role}: {msg['message'][:200]}\n" # Limit to 200 chars per msg
98-
)
133+
# Keep full messages, only truncate very long ones at 500 chars
134+
msg_text = msg['message']
135+
if len(msg_text) > 500:
136+
msg_text = msg_text[:500] + "..."
137+
conversation_context += f"{role}: {msg_text}\n\n"
99138

100139
# Build enhanced prompt with RAG context
101140
system_prompt = None
102141
if rag:
103142
system_prompt = build_system_prompt_with_rag(rag, user_input)
104143
else:
105-
system_prompt = """You are a Clarifai CLI expert. Help users with CLI commands, options, and troubleshooting.
106-
For non-CLI questions, direct users to: https://docs.clarifai.com
107-
108-
IMPORTANT: Keep responses CONCISE and FOCUSED. Use bullet points when appropriate. Maximum 300 words."""
144+
system_prompt = """You are an expert Clarifai CLI assistant. Your role is to help users with:
145+
1. Using the Clarifai CLI commands (login, chat, config, deployment, pipeline, model, artifact, etc.)
146+
2. Understanding CLI options, parameters, and flags
147+
3. Troubleshooting CLI issues and errors
148+
4. Writing CLI scripts and automation
149+
5. Understanding CLI integration with Clarifai resources
150+
6. Answering meta-questions about our conversation
151+
152+
RESPONSE RULES:
153+
- Keep responses CONCISE and FOCUSED (max 300 words)
154+
- Use bullet points, tables, or code examples when helpful
155+
- Answer ALL CLI-related questions directly
156+
- For meta-questions about our conversation, reference the conversation history
157+
- For general Clarifai API questions NOT related to CLI, refer to: https://docs.clarifai.com"""
109158

110159
# Create enhanced input with system prompt, history, and question
111160
enhanced_input = f"{system_prompt}{conversation_context}\n\nCurrent User Question: {user_input}\n\nRespond concisely (max 300 words)."
@@ -131,15 +180,30 @@ def chat(ctx):
131180
{'role': 'assistant', 'message': assistant_message}
132181
)
133182

134-
# Display response
135-
click.secho(f"Assistant: {assistant_message}", fg='green')
183+
# Display response with rich markdown rendering
184+
console.print(
185+
Markdown(f"**Assistant:**\n\n{assistant_message}"),
186+
style='green',
187+
)
136188
else:
137189
click.secho("No text response received", fg='red')
138190
else:
139191
click.secho("Invalid response format", fg='red')
140192

141193
except Exception as e:
142-
click.secho(f"Error calling model: {str(e)}", fg='red')
194+
error_msg = str(e).lower()
195+
196+
# Check for common reasons assistant might be unavailable
197+
if 'not found' in error_msg or 'invalid' in error_msg:
198+
click.secho("Assistant is currently unavailable.", fg='yellow')
199+
click.echo(" Please check your credentials or account status.")
200+
elif 'permission' in error_msg or 'unauthorized' in error_msg:
201+
click.secho("Assistant is currently unavailable.", fg='yellow')
202+
click.echo(" Access denied. Please verify your account status.")
203+
else:
204+
click.secho("Assistant is temporarily unavailable.", fg='yellow')
205+
click.echo(" Please try again in a moment.")
206+
143207
logger.exception("Chat model prediction error")
144208

145209
click.echo() # Add spacing between exchanges
@@ -155,6 +219,9 @@ def chat(ctx):
155219
click.secho(f"Error: Failed to import Clarifai SDK. {str(e)}", fg='red')
156220
sys.exit(1)
157221
except Exception as e:
158-
click.secho(f"Error initializing chat model: {str(e)}", fg='red')
222+
click.secho("Assistant is currently unavailable.\n", fg='yellow')
223+
click.secho("This could be due to:", fg='yellow')
224+
click.echo(" • Network connectivity issues")
225+
click.echo(" • API authentication problems\n")
159226
logger.exception("Chat initialization error")
160227
sys.exit(1)

‎clarifai/cli/rag.py‎

Lines changed: 54 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -10,37 +10,40 @@ class ClarifaiCodeRAG:
1010
def __init__(self, repo_root: str):
1111
"""Initialize RAG with repo root path."""
1212
self.repo_root = Path(repo_root)
13-
self.cli_dir = self.repo_root / "clarifai" / "cli"
14-
self.client_dir = self.repo_root / "clarifai" / "client"
15-
self.utils_dir = self.repo_root / "clarifai" / "utils"
13+
self.clarifai_dir = self.repo_root / "clarifai"
1614
self.documents = []
1715
self._index_codebase()
1816

1917
def _index_codebase(self):
20-
"""Index all Python files in the clarifai package."""
21-
directories = [self.cli_dir, self.client_dir, self.utils_dir]
22-
23-
for directory in directories:
24-
if not directory.exists():
18+
"""Index all Python files in the entire clarifai package."""
19+
# Index the entire clarifai package recursively
20+
if not self.clarifai_dir.exists():
21+
return
22+
23+
for py_file in self.clarifai_dir.rglob("*.py"):
24+
# Skip __pycache__ directories
25+
if "__pycache__" in py_file.parts:
2526
continue
2627

27-
for py_file in directory.rglob("*.py"):
28-
if py_file.name == "__pycache__":
29-
continue
30-
31-
try:
32-
with open(py_file, 'r', encoding='utf-8', errors='ignore') as f:
33-
content = f.read()
34-
relative_path = py_file.relative_to(self.repo_root)
35-
self.documents.append(
36-
{
37-
'path': str(relative_path),
38-
'content': content,
39-
'is_cli': 'cli' in str(relative_path),
40-
}
41-
)
42-
except Exception:
43-
continue
28+
try:
29+
with open(py_file, 'r', encoding='utf-8', errors='ignore') as f:
30+
content = f.read()
31+
relative_path = py_file.relative_to(self.repo_root)
32+
# Prioritize CLI and client modules
33+
module_path = str(relative_path)
34+
is_cli = 'cli' in module_path
35+
is_client = 'client' in module_path
36+
37+
self.documents.append(
38+
{
39+
'path': module_path,
40+
'content': content,
41+
'is_cli': is_cli,
42+
'is_client': is_client,
43+
}
44+
)
45+
except Exception:
46+
continue
4447

4548
def _keyword_match_score(self, query: str, content: str) -> float:
4649
"""Calculate relevance score based on keyword matching."""
@@ -68,21 +71,28 @@ def search(self, query: str, top_k: int = 3) -> List[Tuple[str, str]]:
6871
6972
Returns list of (file_path, relevant_code) tuples.
7073
"""
71-
# Prioritize CLI documents
74+
# Categorize documents by importance for scoring
7275
cli_docs = [d for d in self.documents if d['is_cli']]
73-
other_docs = [d for d in self.documents if not d['is_cli']]
76+
client_docs = [d for d in self.documents if d['is_client'] and not d['is_cli']]
77+
other_docs = [d for d in self.documents if not d['is_cli'] and not d['is_client']]
7478

7579
all_scores = []
7680

81+
# Search all documents with prioritization
7782
for doc in cli_docs:
7883
score = self._keyword_match_score(query, doc['content'])
7984
if score > 0:
8085
all_scores.append((doc, score, 2.0)) # CLI docs get 2x boost
8186

87+
for doc in client_docs:
88+
score = self._keyword_match_score(query, doc['content'])
89+
if score > 0:
90+
all_scores.append((doc, score, 1.5)) # Client docs get 1.5x boost
91+
8292
for doc in other_docs:
8393
score = self._keyword_match_score(query, doc['content'])
8494
if score > 0:
85-
all_scores.append((doc, score, 1.0))
95+
all_scores.append((doc, score, 1.0)) # Other docs (including pipelines, etc)
8696

8797
# Sort by boosted score
8898
all_scores.sort(key=lambda x: x[1] * x[2], reverse=True)
@@ -144,26 +154,27 @@ def build_system_prompt_with_rag(rag: ClarifaiCodeRAG, query: str) -> str:
144154
context += f"\n### From {file_path}:\n```\n{snippet}\n```"
145155

146156
system_prompt = f"""You are an expert Clarifai CLI assistant. Your role is to help users with:
147-
1. Using the Clarifai CLI commands
148-
2. Understanding CLI options and parameters
149-
3. Troubleshooting CLI issues
157+
1. Using the Clarifai CLI commands (login, chat, config, deployment, pipeline, model, artifact, etc.)
158+
2. Understanding CLI options, parameters, and flags
159+
3. Troubleshooting CLI issues and errors
150160
4. Writing CLI scripts and automation
161+
5. Understanding CLI integration with Clarifai resources (pipelines, models, deployments, etc.)
162+
6. Answering meta-questions about our conversation
151163
152-
IMPORTANT RESPONSE RULES:
164+
RESPONSE RULES:
153165
- Keep responses CONCISE and FOCUSED (max 300 words)
154-
- Use bullet points, tables, or short examples when helpful
166+
- Use bullet points, tables, or code examples when helpful
155167
- Get to the point quickly
156-
- ONLY answer questions related to the Clarifai CLI
157-
- For questions about general Clarifai usage (models, apps, data, workflows), respond with:
158-
"I'm specifically designed to help with the Clarifai CLI. For general questions, please refer to the Clarifai documentation at: https://docs.clarifai.com"
159-
- For non-CLI related questions, redirect to: https://docs.clarifai.com
160-
- Always reference relevant CLI files when helpful
161-
- Provide practical examples and command snippets
162-
163-
CLARIFAI CLI CONTEXT:
168+
- Answer ALL CLI-related questions directly (don't redirect)
169+
- For meta-questions about our conversation, reference the conversation history
170+
- For general Clarifai API questions NOT related to CLI, refer to: https://docs.clarifai.com
171+
- Always reference code when available
172+
- Provide practical, working examples
173+
174+
AVAILABLE CLARIFAI CLI COMMANDS:
164175
{rag.get_cli_context()}
165176
166-
CODE REFERENCES FOR THIS QUERY:
167-
{context}"""
177+
RELEVANT CODE SNIPPETS FOR THIS QUESTION:
178+
{context if context else "(Searching full codebase for relevant information)"}"""
168179

169180
return system_prompt

‎requirements.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,3 +18,4 @@ pydantic_core>=2.33.2
1818
packaging>=25.0
1919
tenacity>=8.2.3
2020
httpx>=0.27.0
21+
rich>=13.0.0

0 commit comments

Comments
 (0)