Spaces:
Running
Running
Add Vanna
Browse files- app.py +168 -3
- requirements.txt +4 -0
- src/vanna.py +316 -0
- src/vanna_huggingface_llm_service.py +236 -0
app.py
CHANGED
|
@@ -8,6 +8,7 @@ Now with Datawrapper integration for chart generation!
|
|
| 8 |
"""
|
| 9 |
|
| 10 |
import os
|
|
|
|
| 11 |
import asyncio
|
| 12 |
import pandas as pd
|
| 13 |
import gradio as gr
|
|
@@ -16,6 +17,7 @@ from src.rag_pipeline import create_pipeline
|
|
| 16 |
from src.datawrapper_client import create_and_publish_chart, get_iframe_html
|
| 17 |
from datetime import datetime, timedelta
|
| 18 |
from collections import defaultdict
|
|
|
|
| 19 |
|
| 20 |
# Load environment variables
|
| 21 |
load_dotenv()
|
|
@@ -37,7 +39,20 @@ try:
|
|
| 37 |
except Exception as e:
|
| 38 |
print(f"✗ Error initializing pipeline: {e}")
|
| 39 |
raise
|
| 40 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 41 |
|
| 42 |
def check_rate_limit(request: gr.Request) -> tuple[bool, int]:
|
| 43 |
"""Check if user has exceeded rate limit"""
|
|
@@ -181,6 +196,112 @@ def generate_chart_from_csv(csv_file, user_prompt):
|
|
| 181 |
</div>
|
| 182 |
"""
|
| 183 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 184 |
|
| 185 |
# Minimal CSS to fix UI artifacts and style the mode selector
|
| 186 |
custom_css = """
|
|
@@ -218,6 +339,8 @@ with gr.Blocks(
|
|
| 218 |
with gr.Row():
|
| 219 |
ideation_btn = gr.Button("💡 Ideation Mode", variant="primary", elem_classes="mode-button")
|
| 220 |
chart_gen_btn = gr.Button("📊 Chart Generation Mode", variant="secondary", elem_classes="mode-button")
|
|
|
|
|
|
|
| 221 |
|
| 222 |
# Ideation Mode: Chat interface (shown by default, wrapped in Column)
|
| 223 |
with gr.Column(visible=True) as ideation_container:
|
|
@@ -256,34 +379,67 @@ with gr.Blocks(
|
|
| 256 |
label="Generated Chart"
|
| 257 |
)
|
| 258 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 259 |
# Mode switching functions
|
| 260 |
def switch_to_ideation():
|
| 261 |
return [
|
| 262 |
gr.update(variant="primary"), # ideation_btn
|
| 263 |
gr.update(variant="secondary"), # chart_gen_btn
|
|
|
|
| 264 |
gr.update(visible=True), # ideation_container
|
| 265 |
gr.update(visible=False), # chart_gen_container
|
|
|
|
| 266 |
]
|
| 267 |
|
| 268 |
def switch_to_chart_gen():
|
| 269 |
return [
|
| 270 |
gr.update(variant="secondary"), # ideation_btn
|
| 271 |
gr.update(variant="primary"), # chart_gen_btn
|
|
|
|
| 272 |
gr.update(visible=False), # ideation_container
|
| 273 |
gr.update(visible=True), # chart_gen_container
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 274 |
]
|
| 275 |
|
| 276 |
# Wire up mode switching
|
| 277 |
ideation_btn.click(
|
| 278 |
fn=switch_to_ideation,
|
| 279 |
inputs=[],
|
| 280 |
-
outputs=[ideation_btn, chart_gen_btn, ideation_container, chart_gen_container]
|
| 281 |
)
|
| 282 |
|
| 283 |
chart_gen_btn.click(
|
| 284 |
fn=switch_to_chart_gen,
|
| 285 |
inputs=[],
|
| 286 |
-
outputs=[ideation_btn, chart_gen_btn, ideation_container, chart_gen_container]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 287 |
)
|
| 288 |
|
| 289 |
# Generate chart when button is clicked
|
|
@@ -293,6 +449,13 @@ with gr.Blocks(
|
|
| 293 |
outputs=[chart_output]
|
| 294 |
)
|
| 295 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 296 |
# Knowledge base section (below both interfaces)
|
| 297 |
gr.Markdown("""
|
| 298 |
### About Viz LLM
|
|
@@ -300,6 +463,8 @@ with gr.Blocks(
|
|
| 300 |
**Ideation Mode:** Get design recommendations based on research papers, design principles, and examples from the field of information graphics and data visualization.
|
| 301 |
|
| 302 |
**Chart Generation Mode:** Upload your CSV data and describe your visualization goal. The AI will analyze your data, select the optimal chart type, and generate a publication-ready chart using Datawrapper.
|
|
|
|
|
|
|
| 303 |
|
| 304 |
**Credits:** Special thanks to the researchers whose work informed this model: Robert Kosara, Edward Segel, Jeffrey Heer, Matthew Conlen, John Maeda, Kennedy Elliott, Scott McCloud, and many others.
|
| 305 |
|
|
|
|
| 8 |
"""
|
| 9 |
|
| 10 |
import os
|
| 11 |
+
import io
|
| 12 |
import asyncio
|
| 13 |
import pandas as pd
|
| 14 |
import gradio as gr
|
|
|
|
| 17 |
from src.datawrapper_client import create_and_publish_chart, get_iframe_html
|
| 18 |
from datetime import datetime, timedelta
|
| 19 |
from collections import defaultdict
|
| 20 |
+
from src.vanna import VannaComponent
|
| 21 |
|
| 22 |
# Load environment variables
|
| 23 |
load_dotenv()
|
|
|
|
| 39 |
except Exception as e:
|
| 40 |
print(f"✗ Error initializing pipeline: {e}")
|
| 41 |
raise
|
| 42 |
+
|
| 43 |
+
# Initialize Vanna
|
| 44 |
+
print("Initializing Vanna...")
|
| 45 |
+
try:
|
| 46 |
+
vanna = VannaComponent(
|
| 47 |
+
hf_model="Qwen/Qwen3-VL-30B-A3B-Instruct",
|
| 48 |
+
hf_token=os.getenv("HF_TOKEN_VANNA"),
|
| 49 |
+
hf_provider="novita",
|
| 50 |
+
connection_string=os.getenv("SUPABASE_CONNECTION")
|
| 51 |
+
)
|
| 52 |
+
print("✓ Vanna initialized successfully")
|
| 53 |
+
except Exception as e:
|
| 54 |
+
print(f"✗ Error initializing Vanna: {e}")
|
| 55 |
+
raise
|
| 56 |
|
| 57 |
def check_rate_limit(request: gr.Request) -> tuple[bool, int]:
|
| 58 |
"""Check if user has exceeded rate limit"""
|
|
|
|
| 196 |
</div>
|
| 197 |
"""
|
| 198 |
|
| 199 |
+
def csv_to_cards_html(csv_text: str) -> str:
|
| 200 |
+
"""
|
| 201 |
+
Transforme le CSV brut retourné par Vanna en cartes HTML.
|
| 202 |
+
"""
|
| 203 |
+
try:
|
| 204 |
+
df = pd.read_csv(io.StringIO(csv_text.strip()))
|
| 205 |
+
if df.empty:
|
| 206 |
+
return "<div style='padding: 50px; text-align: center;'>Aucune donnée trouvée.</div>"
|
| 207 |
+
|
| 208 |
+
cards_html = ""
|
| 209 |
+
for _, row in df.iterrows():
|
| 210 |
+
title = row.get("title", "Sans titre")
|
| 211 |
+
source_url = row.get("source_url", "#")
|
| 212 |
+
author = row.get("author", "Inconnu")
|
| 213 |
+
published_date = row.get("published_date", "")
|
| 214 |
+
if not published_date == "nan":
|
| 215 |
+
published_date = ""
|
| 216 |
+
image_url = row.get("image_url", "")
|
| 217 |
+
if not image_url == "nan":
|
| 218 |
+
image_url = "https://fpoimg.com/800x600?text=Image+not+found"
|
| 219 |
+
|
| 220 |
+
cards_html += f"""
|
| 221 |
+
<div style="background: white; border-radius: 10px; box-shadow: 0 2px 8px rgba(0,0,0,0.1);
|
| 222 |
+
overflow: hidden; margin: 10px; width: 320px; flex: 0 0 auto;">
|
| 223 |
+
<img src="{image_url}" alt="{title}" style="width:100%; height:180px; object-fit:cover;">
|
| 224 |
+
<div style="padding: 12px 16px;">
|
| 225 |
+
<h3 style="margin:0; font-size:1.1em; color:#222;">{title}</h3>
|
| 226 |
+
<p style="margin:6px 0; color:#555; font-size:0.9em;">{author}</p>
|
| 227 |
+
<p style="margin:0; color:#999; font-size:0.8em;">{published_date}</p>
|
| 228 |
+
<a href="{source_url}" target="_blank"
|
| 229 |
+
style="display:inline-block; margin-top:8px; font-size:0.9em; color:#1976d2; text-decoration:none;">
|
| 230 |
+
🔗 Voir la source
|
| 231 |
+
</a>
|
| 232 |
+
</div>
|
| 233 |
+
</div>
|
| 234 |
+
"""
|
| 235 |
+
|
| 236 |
+
html = f"""
|
| 237 |
+
<div style="display:flex; flex-wrap:wrap; justify-content:center; padding:20px;">
|
| 238 |
+
{cards_html}
|
| 239 |
+
</div>
|
| 240 |
+
"""
|
| 241 |
+
return html
|
| 242 |
+
|
| 243 |
+
except Exception as e:
|
| 244 |
+
return f"<div style='padding: 50px; text-align: center; color:red;'>Erreur lors du parsing du CSV : {e}</div>"
|
| 245 |
+
|
| 246 |
+
|
| 247 |
+
async def search_inspiration_from_database(user_prompt):
|
| 248 |
+
"""
|
| 249 |
+
Search inspiration posts from user prompt in database.
|
| 250 |
+
|
| 251 |
+
Args:
|
| 252 |
+
user_prompt: User's description of the inspiration query
|
| 253 |
+
|
| 254 |
+
Returns:
|
| 255 |
+
HTML string displaying cards or an error message
|
| 256 |
+
"""
|
| 257 |
+
if not user_prompt or user_prompt.strip() == "":
|
| 258 |
+
return """
|
| 259 |
+
<div style='padding: 50px; text-align: center;'>
|
| 260 |
+
Please describe what kind of inspiration you want to search for.
|
| 261 |
+
</div>
|
| 262 |
+
"""
|
| 263 |
+
|
| 264 |
+
try:
|
| 265 |
+
response = await vanna.ask(user_prompt)
|
| 266 |
+
print("response :", repr(response))
|
| 267 |
+
|
| 268 |
+
clean_response = response.strip()
|
| 269 |
+
|
| 270 |
+
if clean_response.startswith("⚠️") or "Aucun CSV détecté" in clean_response:
|
| 271 |
+
return f"""
|
| 272 |
+
<div style='padding: 50px; text-align: center; color: #d9534f;'>
|
| 273 |
+
<h3>❌ No valid data found</h3>
|
| 274 |
+
<p>The AI couldn't generate any data for this request. Try being more specific — for example:
|
| 275 |
+
<em>"Show me spotlights from 2020 about design"</em>.</p>
|
| 276 |
+
</div>
|
| 277 |
+
"""
|
| 278 |
+
|
| 279 |
+
csv_text = (
|
| 280 |
+
clean_response
|
| 281 |
+
.strip("```")
|
| 282 |
+
.replace("csv", "")
|
| 283 |
+
.replace("CSV", "")
|
| 284 |
+
)
|
| 285 |
+
|
| 286 |
+
if "," not in csv_text:
|
| 287 |
+
return f"""
|
| 288 |
+
<div style='padding: 50px; text-align: center; color: #d9534f;'>
|
| 289 |
+
<h3>❌ No valid CSV detected</h3>
|
| 290 |
+
<p>The model didn't return any structured data. Try rephrasing your query to be more precise.</p>
|
| 291 |
+
</div>
|
| 292 |
+
"""
|
| 293 |
+
|
| 294 |
+
cards_html = csv_to_cards_html(csv_text)
|
| 295 |
+
return cards_html
|
| 296 |
+
|
| 297 |
+
except Exception as e:
|
| 298 |
+
return f"""
|
| 299 |
+
<div style='padding: 50px; text-align: center; color: red;'>
|
| 300 |
+
<h3>❌ Error</h3>
|
| 301 |
+
<p>{str(e)}</p>
|
| 302 |
+
<p style='font-size: 0.9em; color: #666;'>Please try again.</p>
|
| 303 |
+
</div>
|
| 304 |
+
"""
|
| 305 |
|
| 306 |
# Minimal CSS to fix UI artifacts and style the mode selector
|
| 307 |
custom_css = """
|
|
|
|
| 339 |
with gr.Row():
|
| 340 |
ideation_btn = gr.Button("💡 Ideation Mode", variant="primary", elem_classes="mode-button")
|
| 341 |
chart_gen_btn = gr.Button("📊 Chart Generation Mode", variant="secondary", elem_classes="mode-button")
|
| 342 |
+
inspiration_btn = gr.Button("✨ Inspiration Mode", variant="secondary", elem_classes="mode-button")
|
| 343 |
+
|
| 344 |
|
| 345 |
# Ideation Mode: Chat interface (shown by default, wrapped in Column)
|
| 346 |
with gr.Column(visible=True) as ideation_container:
|
|
|
|
| 379 |
label="Generated Chart"
|
| 380 |
)
|
| 381 |
|
| 382 |
+
# Inspiration Mode:
|
| 383 |
+
with gr.Column(visible=False) as inspiration_container:
|
| 384 |
+
with gr.Row():
|
| 385 |
+
inspiration_prompt_input = gr.Textbox(
|
| 386 |
+
placeholder="Ask for an inspiration...",
|
| 387 |
+
show_label=False,
|
| 388 |
+
scale=4,
|
| 389 |
+
container=False
|
| 390 |
+
)
|
| 391 |
+
inspiration_search_btn = gr.Button("🔍 Search", variant="primary", scale=1)
|
| 392 |
+
|
| 393 |
+
inspiration_cards_html = gr.HTML("")
|
| 394 |
+
|
| 395 |
# Mode switching functions
|
| 396 |
def switch_to_ideation():
|
| 397 |
return [
|
| 398 |
gr.update(variant="primary"), # ideation_btn
|
| 399 |
gr.update(variant="secondary"), # chart_gen_btn
|
| 400 |
+
gr.update(variant="secondary"), # inspiration_btn
|
| 401 |
gr.update(visible=True), # ideation_container
|
| 402 |
gr.update(visible=False), # chart_gen_container
|
| 403 |
+
gr.update(visible=False), # inspiration_container
|
| 404 |
]
|
| 405 |
|
| 406 |
def switch_to_chart_gen():
|
| 407 |
return [
|
| 408 |
gr.update(variant="secondary"), # ideation_btn
|
| 409 |
gr.update(variant="primary"), # chart_gen_btn
|
| 410 |
+
gr.update(variant="secondary"), # inspiration_btn
|
| 411 |
gr.update(visible=False), # ideation_container
|
| 412 |
gr.update(visible=True), # chart_gen_container
|
| 413 |
+
gr.update(visible=False), # inspiration_container
|
| 414 |
+
]
|
| 415 |
+
|
| 416 |
+
def switch_to_inspiration():
|
| 417 |
+
return [
|
| 418 |
+
gr.update(variant="secondary"), # ideation_btn
|
| 419 |
+
gr.update(variant="secondary"), # chart_gen_btn
|
| 420 |
+
gr.update(variant="primary"), # inspiration_btn
|
| 421 |
+
gr.update(visible=False), # ideation_container
|
| 422 |
+
gr.update(visible=False), # chart_gen_container
|
| 423 |
+
gr.update(visible=True), # inspiration_container
|
| 424 |
]
|
| 425 |
|
| 426 |
# Wire up mode switching
|
| 427 |
ideation_btn.click(
|
| 428 |
fn=switch_to_ideation,
|
| 429 |
inputs=[],
|
| 430 |
+
outputs=[ideation_btn, chart_gen_btn, inspiration_btn, ideation_container, chart_gen_container, inspiration_container]
|
| 431 |
)
|
| 432 |
|
| 433 |
chart_gen_btn.click(
|
| 434 |
fn=switch_to_chart_gen,
|
| 435 |
inputs=[],
|
| 436 |
+
outputs=[ideation_btn, chart_gen_btn, inspiration_btn, ideation_container, chart_gen_container, inspiration_container]
|
| 437 |
+
)
|
| 438 |
+
|
| 439 |
+
inspiration_btn.click(
|
| 440 |
+
fn=switch_to_inspiration,
|
| 441 |
+
inputs=[],
|
| 442 |
+
outputs=[ideation_btn, chart_gen_btn, inspiration_btn, ideation_container, chart_gen_container, inspiration_container]
|
| 443 |
)
|
| 444 |
|
| 445 |
# Generate chart when button is clicked
|
|
|
|
| 449 |
outputs=[chart_output]
|
| 450 |
)
|
| 451 |
|
| 452 |
+
# Search inspiration when button is clicked
|
| 453 |
+
inspiration_search_btn.click(
|
| 454 |
+
fn=search_inspiration_from_database,
|
| 455 |
+
inputs=[inspiration_prompt_input],
|
| 456 |
+
outputs=[inspiration_cards_html]
|
| 457 |
+
)
|
| 458 |
+
|
| 459 |
# Knowledge base section (below both interfaces)
|
| 460 |
gr.Markdown("""
|
| 461 |
### About Viz LLM
|
|
|
|
| 463 |
**Ideation Mode:** Get design recommendations based on research papers, design principles, and examples from the field of information graphics and data visualization.
|
| 464 |
|
| 465 |
**Chart Generation Mode:** Upload your CSV data and describe your visualization goal. The AI will analyze your data, select the optimal chart type, and generate a publication-ready chart using Datawrapper.
|
| 466 |
+
|
| 467 |
+
**Inspiration Mode:** Coming soon.
|
| 468 |
|
| 469 |
**Credits:** Special thanks to the researchers whose work informed this model: Robert Kosara, Edward Segel, Jeffrey Heer, Matthew Conlen, John Maeda, Kennedy Elliott, Scott McCloud, and many others.
|
| 470 |
|
requirements.txt
CHANGED
|
@@ -17,3 +17,7 @@ pydantic>=2.0.0
|
|
| 17 |
datawrapper>=2.0.7
|
| 18 |
mcp>=1.20.0
|
| 19 |
pandas>=2.0.0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 17 |
datawrapper>=2.0.7
|
| 18 |
mcp>=1.20.0
|
| 19 |
pandas>=2.0.0
|
| 20 |
+
|
| 21 |
+
# Vanna
|
| 22 |
+
requests
|
| 23 |
+
vanna[postgres,chromadb] @ git+https://github.com/vanna-ai/vanna.git@v2
|
src/vanna.py
ADDED
|
@@ -0,0 +1,316 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import re
|
| 2 |
+
import os
|
| 3 |
+
from vanna import Agent, AgentConfig
|
| 4 |
+
from vanna.core.registry import ToolRegistry
|
| 5 |
+
from vanna.core.user import UserResolver, User, RequestContext
|
| 6 |
+
from vanna.tools import RunSqlTool
|
| 7 |
+
from vanna.tools.agent_memory import SaveQuestionToolArgsTool, SearchSavedCorrectToolUsesTool
|
| 8 |
+
from vanna.integrations.postgres import PostgresRunner
|
| 9 |
+
from vanna.integrations.local.agent_memory import DemoAgentMemory
|
| 10 |
+
from .vanna_huggingface_llm_service import VannaHuggingFaceLlmService
|
| 11 |
+
|
| 12 |
+
from typing import List, Dict, Any, Optional
|
| 13 |
+
from vanna.core.system_prompt import SystemPromptBuilder
|
| 14 |
+
from vanna.core.registry import ToolSchema
|
| 15 |
+
from datetime import datetime
|
| 16 |
+
|
| 17 |
+
class CustomSQLSystemPromptBuilder(SystemPromptBuilder):
|
| 18 |
+
"""System prompt builder complet pour SQL assistant Vanna v2."""
|
| 19 |
+
|
| 20 |
+
VERSION = "2.2.0"
|
| 21 |
+
|
| 22 |
+
def __init__(self, company_name: str = "CoJournalist", sql_runner: Optional[PostgresRunner] = None):
|
| 23 |
+
self.company_name = company_name
|
| 24 |
+
self.sql_runner = sql_runner
|
| 25 |
+
|
| 26 |
+
async def build_system_prompt(
|
| 27 |
+
self,
|
| 28 |
+
user: User,
|
| 29 |
+
tool_schemas: List[ToolSchema],
|
| 30 |
+
context: Optional[Dict[str, Any]] = None
|
| 31 |
+
) -> str:
|
| 32 |
+
today = datetime.now().strftime("%Y-%m-%d")
|
| 33 |
+
username = getattr(user, "username", user.id)
|
| 34 |
+
|
| 35 |
+
# ======================
|
| 36 |
+
# BASE DU PROMPT
|
| 37 |
+
# ======================
|
| 38 |
+
prompt = f"[System Prompt v{self.VERSION}]\n\n"
|
| 39 |
+
prompt += f"Vous êtes un assistant SQL expert pour l'entreprise {self.company_name}.\n"
|
| 40 |
+
prompt += f"Date : {today}\nUtilisateur : {username}\nGroupes : {', '.join(user.group_memberships)}\n\n"
|
| 41 |
+
|
| 42 |
+
prompt += (
|
| 43 |
+
"Votre rôle : générer des requêtes SQL correctes et efficaces à partir du langage naturel.\n"
|
| 44 |
+
"Vous répondez toujours au format CSV brut, sans explication ni texte additionnel.\n"
|
| 45 |
+
"Vous avez accès à toutes les tables et relations décrites dans le schéma.\n"
|
| 46 |
+
)
|
| 47 |
+
|
| 48 |
+
# ======================
|
| 49 |
+
# DIRECTIVES SQL
|
| 50 |
+
# ======================
|
| 51 |
+
prompt += (
|
| 52 |
+
"\n## Directives SQL\n"
|
| 53 |
+
"- Toujours utiliser des alias de table dans les JOINs\n"
|
| 54 |
+
"- Ne jamais utiliser SELECT *\n"
|
| 55 |
+
"- Préférer les fonctions fenêtres plutôt que les sous-requêtes quand possible\n"
|
| 56 |
+
"- Toujours inclure un LIMIT pour les requêtes exploratoires\n"
|
| 57 |
+
"- Exclure les posts dont le provider = 'SND'\n"
|
| 58 |
+
"- Exclure les posts dont le type = 'resource'\n"
|
| 59 |
+
"- Exclure les posts dont le type = 'insight'\n"
|
| 60 |
+
"- Formater les dates et nombres pour la lisibilité\n"
|
| 61 |
+
)
|
| 62 |
+
|
| 63 |
+
# ======================
|
| 64 |
+
# SCHÉMA DE LA BASE
|
| 65 |
+
# ======================
|
| 66 |
+
if context and "database_schema" in context:
|
| 67 |
+
prompt += "\n## Schéma de la base de données\n"
|
| 68 |
+
prompt += context["database_schema"]
|
| 69 |
+
else:
|
| 70 |
+
prompt += (
|
| 71 |
+
"\n## Schéma de la base de données\n"
|
| 72 |
+
"Tables :\n"
|
| 73 |
+
"- posts (id, title, source_url, author, published_date, image_url, type, provider_id, created_at, updated_at)\n"
|
| 74 |
+
"- providers (id, name)\n"
|
| 75 |
+
"- provider_attributes (id, provider_id, type, name)\n"
|
| 76 |
+
"- post_provider_attributes (post_id, attribute_id)\n"
|
| 77 |
+
"- tags (id, name)\n"
|
| 78 |
+
"- post_tags (post_id, tag_id, weight)\n"
|
| 79 |
+
"\nRelations :\n"
|
| 80 |
+
" - posts.provider_id → providers.id\n"
|
| 81 |
+
" - post_provider_attributes.post_id → posts.id\n"
|
| 82 |
+
" - post_provider_attributes.attribute_id → provider_attributes.id\n"
|
| 83 |
+
" - provider_attributes.provider_id → providers.id\n"
|
| 84 |
+
" - post_tags.post_id → posts.id\n"
|
| 85 |
+
" - post_tags.tag_id → tags.id\n"
|
| 86 |
+
)
|
| 87 |
+
|
| 88 |
+
# ======================
|
| 89 |
+
# INFORMATIONS SÉMANTIQUES
|
| 90 |
+
# ======================
|
| 91 |
+
prompt += (
|
| 92 |
+
"\n## Informations sémantiques\n"
|
| 93 |
+
"- `posts.title` : titre du contenu (souvent descriptif, peut contenir des mots-clés thématiques).\n"
|
| 94 |
+
"- `posts.source_url` : lien externe vers la ressource ou article.\n"
|
| 95 |
+
"- `posts.author` : nom du journaliste, du média, ou de l’organisation (ex: \"The New York Times\").\n"
|
| 96 |
+
"- `posts.published_date` : date de publication du post.\n"
|
| 97 |
+
"- `posts.type` : type du contenu, de type ENUM ('spotlight', 'resource', 'insight').\n"
|
| 98 |
+
"- `posts.provider_id` : identifiant de la source (provider) ayant publié le contenu.\n"
|
| 99 |
+
"- `providers.name` : nom de l’organisation source (ex: 'Nuanced', 'SND').\n"
|
| 100 |
+
"- `provider_attributes.type` : type d’attribut du provider (ENUM : 'award', 'category').\n"
|
| 101 |
+
"- `provider_attributes.name` : valeur de l’attribut (ex: 'Best Design', 'Investigation').\n"
|
| 102 |
+
"- `tags.name` : thématique ou mot-clé associé au post (ex: '3D', 'AI', 'Investigation').\n"
|
| 103 |
+
"- `post_tags.weight` : poids d’association entre un post et un tag (pertinence).\n"
|
| 104 |
+
)
|
| 105 |
+
|
| 106 |
+
# ======================
|
| 107 |
+
# LOGIQUE SÉMANTIQUE MÉTIER
|
| 108 |
+
# ======================
|
| 109 |
+
prompt += (
|
| 110 |
+
"\n## Logique sémantique métier\n"
|
| 111 |
+
"- Les providers nommés 'SND' représentent des contenus internes à exclure systématiquement.\n"
|
| 112 |
+
"- Une recherche mentionnant une organisation (ex: 'New York Times') doit interroger à la fois `posts.author` et `providers.name`.\n"
|
| 113 |
+
"- Par défaut, les posts à renvoyer sont uniquement ceux de type `spotlight`.\n"
|
| 114 |
+
"- Les posts de type `resource` ou `insight` sont exclus de toutes les requêtes, sauf si l’utilisateur demande explicitement des 'resources'.\n"
|
| 115 |
+
"- Les attributs de type 'award' indiquent des récompenses reçues par le provider.\n"
|
| 116 |
+
"- Les attributs de type 'category' définissent le domaine éditorial (ex: 'Tech', 'Investigation').\n"
|
| 117 |
+
"- Les tags permettent de relier les posts à des thématiques ou disciplines précises.\n"
|
| 118 |
+
"- Un même post peut avoir plusieurs tags, awards ou catégories.\n"
|
| 119 |
+
"- Les requêtes doivent agréger ou joindre les tables en fonction du besoin utilisateur :\n"
|
| 120 |
+
" * Par tag : via `post_tags` et `tags`\n"
|
| 121 |
+
" * Par provider : via `posts.provider_id` → `providers.id`\n"
|
| 122 |
+
" * Par award/category : via `post_provider_attributes` et `provider_attributes`\n"
|
| 123 |
+
"- Si l’utilisateur parle de “posts récents”, filtrer sur `published_date >= CURRENT_DATE - INTERVAL '90 days'`.\n"
|
| 124 |
+
"- Lorsqu'une recherche mentionne un sujet (ex: '3D', 'design', 'AI'), cela correspond à un ou plusieurs `tags.name`.\n"
|
| 125 |
+
"- Lorsqu'une recherche mentionne un auteur ou une organisation, chercher dans `author` et `provider.name`.\n"
|
| 126 |
+
"- Si l'utilisateur mentionne une année (ex: \"en 2021\"), filtrer avec EXTRACT(YEAR FROM published_date) = 2021.\n"
|
| 127 |
+
"- Si l'utilisateur mentionne un mois et une année (ex: \"en mai 2021\"), filtrer avec EXTRACT(MONTH FROM published_date) = 5 ET EXTRACT(YEAR FROM published_date) = 2021.\n"
|
| 128 |
+
"- Si l'utilisateur dit \"récemment\" ou \"dernièrement\", sélectionner les posts des 90 derniers jours.\n"
|
| 129 |
+
"- Si l'utilisateur dit \"cette année\", filtrer avec EXTRACT(YEAR FROM published_date) = EXTRACT(YEAR FROM CURRENT_DATE).\n"
|
| 130 |
+
"- Ne jamais comparer directement published_date à une chaîne comme '2021' ou 'mai 2021'.\n"
|
| 131 |
+
"- Toujours limiter les résultats à 9 lignes maximum pour les requêtes exploratoires.\n"
|
| 132 |
+
)
|
| 133 |
+
|
| 134 |
+
# ======================
|
| 135 |
+
# DÉFINITIONS MÉTIER
|
| 136 |
+
# ======================
|
| 137 |
+
prompt += (
|
| 138 |
+
"\n## Définitions métier\n"
|
| 139 |
+
"- **Post** : post mettant en avant un sujet ou projet particulier.\n"
|
| 140 |
+
"- **Spotlight** : post mettant en avant un sujet ou projet particulier.\n"
|
| 141 |
+
"- **Provider** : entité éditoriale (média, organisation, auteur collectif) responsable du contenu.\n"
|
| 142 |
+
"- **Tag** : thématique, domaine ou mot-clé associé à un post.\n"
|
| 143 |
+
"- **Provider Attribute** : métadonnée du provider (peut être une récompense, une catégorie).\n"
|
| 144 |
+
"- **Award** : distinction reçue par un post.\n"
|
| 145 |
+
"- **Category** : domaine éditorial du post.\n"
|
| 146 |
+
)
|
| 147 |
+
|
| 148 |
+
# ======================
|
| 149 |
+
# OUTILS DISPONIBLES
|
| 150 |
+
# ======================
|
| 151 |
+
if tool_schemas:
|
| 152 |
+
prompt += "\n## Outils disponibles\n"
|
| 153 |
+
for tool in tool_schemas:
|
| 154 |
+
prompt += f"- {tool.name}: {getattr(tool, 'description', 'Pas de description')}\n"
|
| 155 |
+
prompt += f" Paramètres: {getattr(tool, 'parameters', 'N/A')}\n"
|
| 156 |
+
|
| 157 |
+
# ======================
|
| 158 |
+
# MÉMOIRE & CONTEXTE
|
| 159 |
+
# ======================
|
| 160 |
+
tool_names = [t.name for t in tool_schemas]
|
| 161 |
+
has_search = "search_saved_correct_tool_uses" in tool_names
|
| 162 |
+
has_save = "save_question_tool_args" in tool_names
|
| 163 |
+
has_text_memory = "save_text_memory" in tool_names
|
| 164 |
+
|
| 165 |
+
if has_search or has_save or has_text_memory:
|
| 166 |
+
prompt += "\n## Système mémoire\n"
|
| 167 |
+
|
| 168 |
+
if has_search or has_save:
|
| 169 |
+
prompt += "\n• Workflow mémoire des outils :\n"
|
| 170 |
+
if has_search:
|
| 171 |
+
prompt += " - Avant exécution, utilisez search_saved_correct_tool_uses pour détecter les patterns existants.\n"
|
| 172 |
+
if has_save:
|
| 173 |
+
prompt += " - Après succès, utilisez save_question_tool_args pour enregistrer la correspondance.\n"
|
| 174 |
+
|
| 175 |
+
if has_text_memory:
|
| 176 |
+
prompt += "\n• Mémoire textuelle :\n"
|
| 177 |
+
prompt += " - Conservez les schémas, terminologies métier, patterns SQL et préférences utilisateur.\n"
|
| 178 |
+
|
| 179 |
+
# ======================
|
| 180 |
+
# EXEMPLES D’INTERACTIONS
|
| 181 |
+
# ======================
|
| 182 |
+
prompt += (
|
| 183 |
+
"\n## Exemples d'interactions\n"
|
| 184 |
+
"Utilisateur : \"Montre-moi les posts liés à la 3D\"\n"
|
| 185 |
+
"Assistant : [call run_sql with \"SELECT p.id, p.title, p.source_url, p.author, p.published_date, p.image_url, p.type "
|
| 186 |
+
"FROM posts p "
|
| 187 |
+
"JOIN post_tags pt ON p.id = pt.post_id "
|
| 188 |
+
"JOIN tags t ON pt.tag_id = t.id "
|
| 189 |
+
"JOIN providers pr ON p.provider_id = pr.id "
|
| 190 |
+
"WHERE t.name ILIKE '%3D%' AND pr.name != 'SND' AND p.type = 'spotlight' "
|
| 191 |
+
"LIMIT 9;\"]\n"
|
| 192 |
+
"Résultat : \"id,title,source_url,author,published_date,image_url,type\"\n"
|
| 193 |
+
"\nUtilisateur : \"Montre-moi les posts du New York Times\"\n"
|
| 194 |
+
"Assistant : [call run_sql with \"SELECT p.id, p.title, p.source_url, p.author, p.published_date, p.image_url, p.type "
|
| 195 |
+
"FROM posts p "
|
| 196 |
+
"LEFT JOIN providers pr ON pr.id = p.provider_id "
|
| 197 |
+
"WHERE LOWER(p.author) LIKE '%new york times%' OR LOWER(pr.name) LIKE '%new york times%' "
|
| 198 |
+
"AND pr.name != 'SND' AND p.type = 'spotlight' "
|
| 199 |
+
"LIMIT 9;\"]\n"
|
| 200 |
+
"Résultat : \"id,title,source_url,author,published_date,image_url\"\n"
|
| 201 |
+
)
|
| 202 |
+
|
| 203 |
+
# ======================
|
| 204 |
+
# INSTRUCTIONS FINALES
|
| 205 |
+
# ======================
|
| 206 |
+
prompt += (
|
| 207 |
+
"\nIMPORTANT :\n"
|
| 208 |
+
"- Toujours exclure les posts dont provider = 'SND'.\n"
|
| 209 |
+
"- Toujours exclure les posts dont type = 'resource'.\n"
|
| 210 |
+
"- Toujours exclure les posts dont type = 'insight'.\n"
|
| 211 |
+
"- Toujours renvoyer uniquement le résultat brut CSV, sans texte ni commentaires.\n"
|
| 212 |
+
"- Ne pas inclure de JSON, d’analyse, ni de messages explicatifs.\n"
|
| 213 |
+
"- Ignorer les itérations supplémentaires ou réflexions internes.\n"
|
| 214 |
+
"- Une fois le résultat obtenu, arrêtez l’exécution du tool.\n"
|
| 215 |
+
)
|
| 216 |
+
|
| 217 |
+
return prompt
|
| 218 |
+
|
| 219 |
+
|
| 220 |
+
class SimpleUserResolver(UserResolver):
|
| 221 |
+
async def resolve_user(self, request_context: RequestContext) -> User:
|
| 222 |
+
user_email = request_context.get_cookie('vanna_email') or 'guest@example.com'
|
| 223 |
+
group = 'admin' if user_email == 'admin@example.com' else 'user'
|
| 224 |
+
return User(id=user_email, email=user_email, group_memberships=[group])
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
class VannaComponent:
|
| 228 |
+
def __init__(
|
| 229 |
+
self,
|
| 230 |
+
hf_model: str,
|
| 231 |
+
hf_token: str,
|
| 232 |
+
hf_provider: str,
|
| 233 |
+
connection_string: str,
|
| 234 |
+
):
|
| 235 |
+
# Configure LLM
|
| 236 |
+
llm = VannaHuggingFaceLlmService(model=hf_model, token=hf_token, provider=hf_provider)
|
| 237 |
+
|
| 238 |
+
# Configure database tool
|
| 239 |
+
self.sql_runner = PostgresRunner(connection_string=connection_string)
|
| 240 |
+
db_tool = RunSqlTool(
|
| 241 |
+
sql_runner=self.sql_runner,
|
| 242 |
+
)
|
| 243 |
+
|
| 244 |
+
# Configure agent memory
|
| 245 |
+
agent_memory = DemoAgentMemory(max_items=1000)
|
| 246 |
+
save_memory_tool = SaveQuestionToolArgsTool(agent_memory)
|
| 247 |
+
search_memory_tool = SearchSavedCorrectToolUsesTool(agent_memory)
|
| 248 |
+
|
| 249 |
+
# Configure user resolver
|
| 250 |
+
self.user_resolver = SimpleUserResolver()
|
| 251 |
+
|
| 252 |
+
# Register tools with access control
|
| 253 |
+
tools = ToolRegistry()
|
| 254 |
+
tools.register_local_tool(db_tool, access_groups=['admin', 'user'])
|
| 255 |
+
tools.register_local_tool(save_memory_tool, access_groups=['admin'])
|
| 256 |
+
tools.register_local_tool(search_memory_tool, access_groups=['admin', 'user'])
|
| 257 |
+
|
| 258 |
+
# Create the agent
|
| 259 |
+
self.agent = Agent(
|
| 260 |
+
llm_service=llm,
|
| 261 |
+
tool_registry=tools,
|
| 262 |
+
user_resolver=self.user_resolver,
|
| 263 |
+
system_prompt_builder=CustomSQLSystemPromptBuilder("CoJournalist", self.sql_runner),
|
| 264 |
+
config=AgentConfig(stream_responses=False, max_tool_iterations=1)
|
| 265 |
+
)
|
| 266 |
+
|
| 267 |
+
async def ask(self, prompt_for_llm: str):
|
| 268 |
+
ctx = RequestContext()
|
| 269 |
+
|
| 270 |
+
print(f"🙋 Prompt envoyé au LLM : {prompt_for_llm}")
|
| 271 |
+
|
| 272 |
+
final_text = ""
|
| 273 |
+
seen_texts = set()
|
| 274 |
+
# 🔁 Boucle sur les composants produits par l'agent
|
| 275 |
+
async for component in self.agent.send_message(request_context=ctx, message=prompt_for_llm):
|
| 276 |
+
# Texte simple produit par l'agent
|
| 277 |
+
simple = getattr(component, "simple_component", None)
|
| 278 |
+
text = getattr(simple, "text", "") if simple else ""
|
| 279 |
+
if text and text not in seen_texts:
|
| 280 |
+
print(f"💬 LLM says (part): {text[:200]}...")
|
| 281 |
+
final_text += text + "\n"
|
| 282 |
+
seen_texts.add(text)
|
| 283 |
+
|
| 284 |
+
# Requête SQL générée (si présente)
|
| 285 |
+
sql_query = getattr(component, "sql", None)
|
| 286 |
+
if sql_query:
|
| 287 |
+
print(f"🧾 Requête SQL générée : {sql_query}")
|
| 288 |
+
|
| 289 |
+
# Métadonnées et autres infos associées au composant
|
| 290 |
+
metadata = getattr(component, "metadata", None)
|
| 291 |
+
if metadata:
|
| 292 |
+
print(f"📋 Métadonnées composant : {metadata}")
|
| 293 |
+
|
| 294 |
+
# Type de composant utile pour debug
|
| 295 |
+
component_type = getattr(component, "type", None)
|
| 296 |
+
if component_type:
|
| 297 |
+
print(f"🔖 Type composant : {component_type}")
|
| 298 |
+
|
| 299 |
+
|
| 300 |
+
match = re.search(r"query_results_[\w-]+\.csv", final_text)
|
| 301 |
+
if match:
|
| 302 |
+
filename = match.group(0)
|
| 303 |
+
folder = "513935c4d2db2d2d"
|
| 304 |
+
full_path = os.path.join(folder, filename)
|
| 305 |
+
|
| 306 |
+
if os.path.exists(full_path):
|
| 307 |
+
print(f"📂 Lecture directe du fichier CSV : {full_path}")
|
| 308 |
+
with open(full_path, "r", encoding="utf-8") as f:
|
| 309 |
+
csv_data = f.read().strip()
|
| 310 |
+
print("🤖 Réponse envoyée à l'user (depuis fichier) :", csv_data[:300])
|
| 311 |
+
return csv_data
|
| 312 |
+
else:
|
| 313 |
+
print(f"⚠️ Fichier non trouvé : {full_path}")
|
| 314 |
+
|
| 315 |
+
return final_text
|
| 316 |
+
|
src/vanna_huggingface_llm_service.py
ADDED
|
@@ -0,0 +1,236 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
from typing import Any, AsyncGenerator, Dict, List, Optional
|
| 5 |
+
|
| 6 |
+
from vanna.core.llm import (
|
| 7 |
+
LlmService,
|
| 8 |
+
LlmRequest,
|
| 9 |
+
LlmResponse,
|
| 10 |
+
LlmStreamChunk,
|
| 11 |
+
)
|
| 12 |
+
from vanna.core.tool import ToolCall, ToolSchema
|
| 13 |
+
from huggingface_hub import InferenceClient
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class VannaHuggingFaceLlmService(LlmService):
|
| 17 |
+
def __init__(
|
| 18 |
+
self,
|
| 19 |
+
model: Optional[str] = None,
|
| 20 |
+
api_key: Optional[str] = None,
|
| 21 |
+
provider: Optional[str] = None,
|
| 22 |
+
**extra_client_kwargs: Any,
|
| 23 |
+
) -> None:
|
| 24 |
+
|
| 25 |
+
"""Initialise le client Hugging Face InferenceClient."""
|
| 26 |
+
client_kwargs = extra_client_kwargs.copy()
|
| 27 |
+
if model:
|
| 28 |
+
client_kwargs["model"] = model
|
| 29 |
+
if api_key:
|
| 30 |
+
client_kwargs["api_key"] = api_key
|
| 31 |
+
if provider:
|
| 32 |
+
client_kwargs["provider"] = provider
|
| 33 |
+
|
| 34 |
+
self.model = model
|
| 35 |
+
|
| 36 |
+
self._client = InferenceClient(**client_kwargs)
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
async def send_request(self, request: LlmRequest) -> LlmResponse:
|
| 40 |
+
"""Send a non-streaming request to OpenAI and return the response."""
|
| 41 |
+
payload = self._build_payload(request)
|
| 42 |
+
|
| 43 |
+
# Call the API synchronously; this function is async but we can block here.
|
| 44 |
+
resp = self._client.chat.completions.create(**payload, stream=False)
|
| 45 |
+
|
| 46 |
+
if not resp.choices:
|
| 47 |
+
return LlmResponse(content=None, tool_calls=None, finish_reason=None)
|
| 48 |
+
|
| 49 |
+
choice = resp.choices[0]
|
| 50 |
+
content: Optional[str] = getattr(choice.message, "content", None)
|
| 51 |
+
tool_calls = self._extract_tool_calls_from_message(choice.message)
|
| 52 |
+
|
| 53 |
+
usage: Dict[str, int] = {}
|
| 54 |
+
if getattr(resp, "usage", None):
|
| 55 |
+
usage = {
|
| 56 |
+
k: int(v)
|
| 57 |
+
for k, v in {
|
| 58 |
+
"prompt_tokens": getattr(resp.usage, "prompt_tokens", 0),
|
| 59 |
+
"completion_tokens": getattr(resp.usage, "completion_tokens", 0),
|
| 60 |
+
"total_tokens": getattr(resp.usage, "total_tokens", 0),
|
| 61 |
+
}.items()
|
| 62 |
+
}
|
| 63 |
+
|
| 64 |
+
return LlmResponse(
|
| 65 |
+
content=content,
|
| 66 |
+
tool_calls=tool_calls or None,
|
| 67 |
+
finish_reason=getattr(choice, "finish_reason", None),
|
| 68 |
+
usage=usage or None,
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
async def stream_request(
|
| 72 |
+
self, request: LlmRequest
|
| 73 |
+
) -> AsyncGenerator[LlmStreamChunk, None]:
|
| 74 |
+
"""Stream a request to OpenAI.
|
| 75 |
+
|
| 76 |
+
Emits `LlmStreamChunk` for textual deltas as they arrive. Tool-calls are
|
| 77 |
+
accumulated and emitted in a final chunk when the stream ends.
|
| 78 |
+
"""
|
| 79 |
+
payload = self._build_payload(request)
|
| 80 |
+
|
| 81 |
+
# Synchronous streaming iterator; iterate within async context.
|
| 82 |
+
stream = self._client.chat.completions.create(**payload, stream=True)
|
| 83 |
+
|
| 84 |
+
# Builders for streamed tool-calls (index -> partial)
|
| 85 |
+
tc_builders: Dict[int, Dict[str, Optional[str]]] = {}
|
| 86 |
+
last_finish: Optional[str] = None
|
| 87 |
+
|
| 88 |
+
for event in stream:
|
| 89 |
+
if not getattr(event, "choices", None):
|
| 90 |
+
continue
|
| 91 |
+
|
| 92 |
+
choice = event.choices[0]
|
| 93 |
+
delta = getattr(choice, "delta", None)
|
| 94 |
+
if delta is None:
|
| 95 |
+
# Some SDK versions use `event.choices[0].message` on the final packet
|
| 96 |
+
last_finish = getattr(choice, "finish_reason", last_finish)
|
| 97 |
+
continue
|
| 98 |
+
|
| 99 |
+
# Text content
|
| 100 |
+
content_piece: Optional[str] = getattr(delta, "content", None)
|
| 101 |
+
if content_piece:
|
| 102 |
+
yield LlmStreamChunk(content=content_piece)
|
| 103 |
+
|
| 104 |
+
# Tool calls (streamed)
|
| 105 |
+
streamed_tool_calls = getattr(delta, "tool_calls", None)
|
| 106 |
+
if streamed_tool_calls:
|
| 107 |
+
for tc in streamed_tool_calls:
|
| 108 |
+
idx = getattr(tc, "index", 0) or 0
|
| 109 |
+
b = tc_builders.setdefault(
|
| 110 |
+
idx, {"id": None, "name": None, "arguments": ""}
|
| 111 |
+
)
|
| 112 |
+
if getattr(tc, "id", None):
|
| 113 |
+
b["id"] = tc.id
|
| 114 |
+
fn = getattr(tc, "function", None)
|
| 115 |
+
if fn is not None:
|
| 116 |
+
if getattr(fn, "name", None):
|
| 117 |
+
b["name"] = fn.name
|
| 118 |
+
if getattr(fn, "arguments", None):
|
| 119 |
+
b["arguments"] = (b["arguments"] or "") + fn.arguments
|
| 120 |
+
|
| 121 |
+
last_finish = getattr(choice, "finish_reason", last_finish)
|
| 122 |
+
|
| 123 |
+
# Emit final tool-calls chunk if any
|
| 124 |
+
final_tool_calls: List[ToolCall] = []
|
| 125 |
+
for b in tc_builders.values():
|
| 126 |
+
if not b.get("name"):
|
| 127 |
+
continue
|
| 128 |
+
args_raw = b.get("arguments") or "{}"
|
| 129 |
+
try:
|
| 130 |
+
loaded = json.loads(args_raw)
|
| 131 |
+
if isinstance(loaded, dict):
|
| 132 |
+
args_dict: Dict[str, Any] = loaded
|
| 133 |
+
else:
|
| 134 |
+
args_dict = {"args": loaded}
|
| 135 |
+
except Exception:
|
| 136 |
+
args_dict = {"_raw": args_raw}
|
| 137 |
+
final_tool_calls.append(
|
| 138 |
+
ToolCall(
|
| 139 |
+
id=b.get("id") or "tool_call",
|
| 140 |
+
name=b["name"] or "tool",
|
| 141 |
+
arguments=args_dict,
|
| 142 |
+
)
|
| 143 |
+
)
|
| 144 |
+
|
| 145 |
+
if final_tool_calls:
|
| 146 |
+
yield LlmStreamChunk(tool_calls=final_tool_calls, finish_reason=last_finish)
|
| 147 |
+
else:
|
| 148 |
+
# Still emit a terminal chunk to signal completion
|
| 149 |
+
yield LlmStreamChunk(finish_reason=last_finish or "stop")
|
| 150 |
+
|
| 151 |
+
async def validate_tools(self, tools: List[ToolSchema]) -> List[str]:
|
| 152 |
+
"""Validate tool schemas. Returns a list of error messages."""
|
| 153 |
+
errors: List[str] = []
|
| 154 |
+
# Basic checks; OpenAI will enforce further validation server-side.
|
| 155 |
+
for t in tools:
|
| 156 |
+
if not t.name or len(t.name) > 64:
|
| 157 |
+
errors.append(f"Invalid tool name: {t.name!r}")
|
| 158 |
+
return errors
|
| 159 |
+
|
| 160 |
+
# Internal helpers
|
| 161 |
+
def _build_payload(self, request: LlmRequest) -> Dict[str, Any]:
|
| 162 |
+
messages: List[Dict[str, Any]] = []
|
| 163 |
+
|
| 164 |
+
# Add system prompt as first message if provided
|
| 165 |
+
if request.system_prompt:
|
| 166 |
+
messages.append({"role": "system", "content": request.system_prompt})
|
| 167 |
+
|
| 168 |
+
for m in request.messages:
|
| 169 |
+
msg: Dict[str, Any] = {"role": m.role, "content": m.content}
|
| 170 |
+
if m.role == "tool" and m.tool_call_id:
|
| 171 |
+
msg["tool_call_id"] = m.tool_call_id
|
| 172 |
+
elif m.role == "assistant" and m.tool_calls:
|
| 173 |
+
# Convert tool calls to OpenAI format
|
| 174 |
+
tool_calls_payload = []
|
| 175 |
+
for tc in m.tool_calls:
|
| 176 |
+
tool_calls_payload.append({
|
| 177 |
+
"id": tc.id,
|
| 178 |
+
"type": "function",
|
| 179 |
+
"function": {
|
| 180 |
+
"name": tc.name,
|
| 181 |
+
"arguments": json.dumps(tc.arguments)
|
| 182 |
+
}
|
| 183 |
+
})
|
| 184 |
+
msg["tool_calls"] = tool_calls_payload
|
| 185 |
+
messages.append(msg)
|
| 186 |
+
|
| 187 |
+
tools_payload: Optional[List[Dict[str, Any]]] = None
|
| 188 |
+
if request.tools:
|
| 189 |
+
tools_payload = [
|
| 190 |
+
{
|
| 191 |
+
"type": "function",
|
| 192 |
+
"function": {
|
| 193 |
+
"name": t.name,
|
| 194 |
+
"description": t.description,
|
| 195 |
+
"parameters": t.parameters,
|
| 196 |
+
},
|
| 197 |
+
}
|
| 198 |
+
for t in request.tools
|
| 199 |
+
]
|
| 200 |
+
|
| 201 |
+
payload: Dict[str, Any] = {
|
| 202 |
+
"model": self.model,
|
| 203 |
+
"messages": messages,
|
| 204 |
+
}
|
| 205 |
+
if request.max_tokens is not None:
|
| 206 |
+
payload["max_tokens"] = request.max_tokens
|
| 207 |
+
if tools_payload:
|
| 208 |
+
payload["tools"] = tools_payload
|
| 209 |
+
payload["tool_choice"] = "auto"
|
| 210 |
+
|
| 211 |
+
return payload
|
| 212 |
+
|
| 213 |
+
def _extract_tool_calls_from_message(self, message: Any) -> List[ToolCall]:
|
| 214 |
+
tool_calls: List[ToolCall] = []
|
| 215 |
+
raw_tool_calls = getattr(message, "tool_calls", None) or []
|
| 216 |
+
for tc in raw_tool_calls:
|
| 217 |
+
fn = getattr(tc, "function", None)
|
| 218 |
+
if not fn:
|
| 219 |
+
continue
|
| 220 |
+
args_raw = getattr(fn, "arguments", "{}")
|
| 221 |
+
try:
|
| 222 |
+
loaded = json.loads(args_raw)
|
| 223 |
+
if isinstance(loaded, dict):
|
| 224 |
+
args_dict: Dict[str, Any] = loaded
|
| 225 |
+
else:
|
| 226 |
+
args_dict = {"args": loaded}
|
| 227 |
+
except Exception:
|
| 228 |
+
args_dict = {"_raw": args_raw}
|
| 229 |
+
tool_calls.append(
|
| 230 |
+
ToolCall(
|
| 231 |
+
id=getattr(tc, "id", "tool_call"),
|
| 232 |
+
name=getattr(fn, "name", "tool"),
|
| 233 |
+
arguments=args_dict,
|
| 234 |
+
)
|
| 235 |
+
)
|
| 236 |
+
return tool_calls
|