remdms commited on
Commit
9db289b
·
1 Parent(s): 7114af0
Files changed (4) hide show
  1. app.py +168 -3
  2. requirements.txt +4 -0
  3. src/vanna.py +316 -0
  4. 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