Preparing proper gemini integrations for the api handlers. Bridging database building gaps in the flow.

This commit is contained in:
Christbru 2025-10-19 07:22:13 -05:00
commit 9ff012dd1d
5 changed files with 131 additions and 56 deletions

View file

@ -1,9 +1,8 @@
use crate::gemini_client;
use crate::vector_db::QdrantClient; use crate::vector_db::QdrantClient;
use crate::storage; use crate::storage;
use anyhow::Result; use anyhow::Result;
use bytes::Buf; use bytes::Buf;
use futures_util::{StreamExt, TryStreamExt}; use futures_util::TryStreamExt;
use serde::Deserialize; use serde::Deserialize;
use sqlx::{MySqlPool, Row}; use sqlx::{MySqlPool, Row};
use warp::{multipart::FormData, Filter, Rejection, Reply}; use warp::{multipart::FormData, Filter, Rejection, Reply};
@ -80,10 +79,7 @@ pub fn routes(pool: MySqlPool) -> impl Filter<Extract = impl Reply, Error = Reje
} }
async fn handle_upload(mut form: FormData, pool: MySqlPool) -> Result<impl Reply, Rejection> { async fn handle_upload(mut form: FormData, pool: MySqlPool) -> Result<impl Reply, Rejection> {
// qdrant client let mut created_files = Vec::new();
let qdrant_url = std::env::var("QDRANT_URL").unwrap_or_else(|_| "http://qdrant:6333".to_string());
let qdrant = QdrantClient::new(&qdrant_url);
while let Some(field) = form.try_next().await.map_err(|_| warp::reject())? { while let Some(field) = form.try_next().await.map_err(|_| warp::reject())? {
let _name = field.name().to_string(); let _name = field.name().to_string();
let filename = field let filename = field
@ -116,7 +112,7 @@ async fn handle_upload(mut form: FormData, pool: MySqlPool) -> Result<impl Reply
// Insert file record with pending_analysis = true, description = NULL // Insert file record with pending_analysis = true, description = NULL
let id = uuid::Uuid::new_v4().to_string(); let id = uuid::Uuid::new_v4().to_string();
sqlx::query("INSERT INTO files (id, filename, path, description, pending_analysis) VALUES (?, ?, ?, ?, ?)") sqlx::query("INSERT INTO files (id, filename, path, description, pending_analysis, analysis_status) VALUES (?, ?, ?, ?, ?, 'Queued')")
.bind(&id) .bind(&id)
.bind(&filename) .bind(&filename)
.bind(path.to_str().unwrap()) .bind(path.to_str().unwrap())
@ -128,10 +124,18 @@ async fn handle_upload(mut form: FormData, pool: MySqlPool) -> Result<impl Reply
tracing::error!("DB insert error: {}", e); tracing::error!("DB insert error: {}", e);
warp::reject() warp::reject()
})?; })?;
// Enqueue worker task to process file (to be implemented) created_files.push(serde_json::json!({
"id": id,
"filename": filename,
"pending_analysis": true,
"analysis_status": "Queued"
}));
} }
Ok(warp::reply::json(&serde_json::json!({"success": true}))) Ok(warp::reply::json(&serde_json::json!({
"uploaded": created_files.len(),
"files": created_files
})))
} }
async fn handle_import_demo(params: std::collections::HashMap<String, String>, pool: MySqlPool) -> Result<impl Reply, Rejection> { async fn handle_import_demo(params: std::collections::HashMap<String, String>, pool: MySqlPool) -> Result<impl Reply, Rejection> {
@ -209,7 +213,7 @@ async fn handle_delete(q: DeleteQuery, pool: MySqlPool) -> Result<impl Reply, Re
let _ = storage::delete_file(std::path::Path::new(&path)); let _ = storage::delete_file(std::path::Path::new(&path));
// Remove from Qdrant // Remove from Qdrant
let qdrant_url = std::env::var("QDRANT_URL").unwrap_or_else(|_| "http://qdrant:6333".to_string()); let qdrant_url = std::env::var("QDRANT_URL").unwrap_or_else(|_| "http://qdrant:6333".to_string());
let qdrant = crate::vector_db::QdrantClient::new(&qdrant_url); let qdrant = QdrantClient::new(&qdrant_url);
let _ = qdrant.delete_point(&q.id).await; let _ = qdrant.delete_point(&q.id).await;
let _ = sqlx::query("DELETE FROM files WHERE id = ?").bind(&q.id).execute(&pool).await; let _ = sqlx::query("DELETE FROM files WHERE id = ?").bind(&q.id).execute(&pool).await;
return Ok(warp::reply::json(&serde_json::json!({"deleted": true}))); return Ok(warp::reply::json(&serde_json::json!({"deleted": true})));

View file

@ -1,4 +1,5 @@
use crate::gemini_client::{generate_text, demo_text_embedding, DEMO_EMBED_DIM}; use crate::gemini_client::{demo_text_embedding, generate_text_with_model, DEMO_EMBED_DIM};
use crate::vector;
use crate::vector_db::QdrantClient; use crate::vector_db::QdrantClient;
use sqlx::MySqlPool; use sqlx::MySqlPool;
use anyhow::Result; use anyhow::Result;
@ -27,6 +28,9 @@ impl FileWorker {
info!("Processing file {}", fid); info!("Processing file {}", fid);
if let Err(e) = self.process_file(&fid).await { if let Err(e) = self.process_file(&fid).await {
error!("Error processing file {}: {}", fid, e); error!("Error processing file {}: {}", fid, e);
if let Err(mark_err) = self.mark_failed(&fid, &format!("{}", e)).await {
error!("Failed to mark file {} as failed: {}", fid, mark_err);
}
} }
} }
Ok(None) => { Ok(None) => {
@ -66,12 +70,18 @@ impl FileWorker {
.bind(file_id) .bind(file_id)
.fetch_one(&self.pool) .fetch_one(&self.pool)
.await?; .await?;
let filename: String = row.get("filename"); let filename: String = row.get("filename");
let path: String = row.get("path"); let _path: String = row.get("path");
// Stage 1: Gemini 2.5 Flash for description // Stage 1: Gemini 2.5 Flash for description
std::env::set_var("GEMINI_MODEL", "gemini-1.5-flash"); let desc = generate_text_with_model(
let desc = generate_text(&format!("Describe the file '{filename}' and extract all key components, keywords, and details for later vectorization. Be comprehensive and factual.")).await.unwrap_or_else(|e| format!("[desc error: {}]", e)); "gemini-2.5-flash",
&format!(
"Describe the file '{filename}' and extract all key components, keywords, and details for later vectorization. Be comprehensive and factual."
),
)
.await
.unwrap_or_else(|e| format!("[desc error: {}]", e));
sqlx::query("UPDATE files SET description = ?, analysis_status = 'InProgress' WHERE id = ?") sqlx::query("UPDATE files SET description = ?, analysis_status = 'InProgress' WHERE id = ?")
.bind(&desc) .bind(&desc)
.bind(file_id) .bind(file_id)
@ -79,12 +89,26 @@ impl FileWorker {
.await?; .await?;
// Stage 2: Gemini 2.5 Pro for deep vector graph data // Stage 2: Gemini 2.5 Pro for deep vector graph data
std::env::set_var("GEMINI_MODEL", "gemini-1.5-pro"); let vector_graph = generate_text_with_model(
let vector_graph = generate_text(&format!("Given the file '{filename}' and its description: {desc}\nGenerate a set of vector graph data (keywords, use cases, relationships) that can be used for broad and precise search. Only include what is directly supported by the file.")).await.unwrap_or_else(|e| format!("[vector error: {}]", e)); "gemini-2.5-pro",
&format!(
"Given the file '{filename}' and its description: {desc}\nGenerate a set of vector graph data (keywords, use cases, relationships) that can be used for broad and precise search. Only include what is directly supported by the file."
),
)
.await
.unwrap_or_else(|e| format!("[vector error: {}]", e));
// Stage 3: Embed and upsert to Qdrant // Stage 3: Embed and upsert to Qdrant
let emb = demo_text_embedding(&vector_graph).await?; let emb = demo_text_embedding(&vector_graph).await?;
self.qdrant.upsert_point(file_id, emb).await?; match self.qdrant.upsert_point(file_id, emb.clone()).await {
Ok(_) => {
let _ = vector::store_embedding(file_id, emb.clone());
}
Err(err) => {
error!("Qdrant upsert failed for {}: {}", file_id, err);
let _ = vector::store_embedding(file_id, emb);
}
}
// Mark file as ready // Mark file as ready
sqlx::query("UPDATE files SET pending_analysis = FALSE, analysis_status = 'Completed' WHERE id = ?") sqlx::query("UPDATE files SET pending_analysis = FALSE, analysis_status = 'Completed' WHERE id = ?")
@ -93,4 +117,17 @@ impl FileWorker {
.await?; .await?;
Ok(()) Ok(())
} }
async fn mark_failed(&self, file_id: &str, reason: &str) -> Result<()> {
sqlx::query("UPDATE files SET analysis_status = 'Failed', pending_analysis = TRUE WHERE id = ?")
.bind(file_id)
.execute(&self.pool)
.await?;
sqlx::query("UPDATE files SET description = COALESCE(description, ?) WHERE id = ?")
.bind(format!("[analysis failed: {}]", reason))
.bind(file_id)
.execute(&self.pool)
.await?;
Ok(())
}
} }

View file

@ -1,30 +1,11 @@
use anyhow::Result; use anyhow::Result;
use serde::{Deserialize, Serialize};
use serde_json::json;
use reqwest::Client; use reqwest::Client;
use serde::Deserialize;
use serde_json::json;
// NOTE: This is a small stub to represent where you'd call the Gemini API. // NOTE: This file provides lightweight helpers around the Gemini API. For the
// Replace with real API call and proper auth handling for production. // hackathon demo we fall back to deterministic strings when the API key is not
// configured so the flows still work end-to-end.
#[derive(Debug, Deserialize)]
pub struct GeminiTokenResponse {
pub token: String,
}
pub async fn generate_token_for_file(_path: &str) -> Result<String> {
Ok("gemini-token-placeholder".to_string())
}
/// Demo embedding generator - deterministic pseudo-embedding from filename/path
pub fn demo_embedding_from_path(path: &str) -> Vec<f32> {
// Very simple: hash bytes into a small vector
let mut v = vec![0f32; 64];
for (i, b) in path.as_bytes().iter().enumerate() {
let idx = i % v.len();
v[idx] += (*b as f32) / 255.0;
}
v
}
pub const DEMO_EMBED_DIM: usize = 64; pub const DEMO_EMBED_DIM: usize = 64;
@ -38,16 +19,27 @@ pub async fn demo_text_embedding(text: &str) -> Result<Vec<f32>> {
Ok(v) Ok(v)
} }
/// Generate text with Gemini (Generative Language API). Falls back to a demo string if GEMINI_API_KEY is not set. /// Generate text using the default model (GEMINI_MODEL or gemini-2.5-pro).
#[allow(dead_code)]
pub async fn generate_text(prompt: &str) -> Result<String> { pub async fn generate_text(prompt: &str) -> Result<String> {
let model = std::env::var("GEMINI_MODEL").unwrap_or_else(|_| "gemini-2.5-pro".to_string());
generate_text_with_model(&model, prompt).await
}
/// Generate text with an explicit Gemini model. Falls back to a deterministic
/// response when the API key is not set so the demo still runs.
pub async fn generate_text_with_model(model: &str, prompt: &str) -> Result<String> {
let api_key = match std::env::var("GEMINI_API_KEY") { let api_key = match std::env::var("GEMINI_API_KEY") {
Ok(k) if !k.is_empty() => k, Ok(k) if !k.is_empty() => k,
_ => { _ => {
return Ok(format!("[demo] Gemini not configured. Prompt preview: {}", truncate(prompt, 240))); return Ok(format!(
"[demo] Gemini ({}) not configured. Prompt preview: {}",
model,
truncate(prompt, 240)
));
} }
}; };
let model = std::env::var("GEMINI_MODEL").unwrap_or_else(|_| "gemini-1.5-pro".to_string());
let url = format!( let url = format!(
"https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent?key={}", "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent?key={}",
model, api_key model, api_key
@ -62,7 +54,12 @@ pub async fn generate_text(prompt: &str) -> Result<String> {
let status = resp.status(); let status = resp.status();
let txt = resp.text().await?; let txt = resp.text().await?;
if !status.is_success() { if !status.is_success() {
return Ok(format!("[demo] Gemini error {}: {}", status, truncate(&txt, 240))); return Ok(format!(
"[demo] Gemini ({}) error {}: {}",
model,
status,
truncate(&txt, 240)
));
} }
#[derive(Deserialize)] #[derive(Deserialize)]
@ -84,5 +81,9 @@ pub async fn generate_text(prompt: &str) -> Result<String> {
} }
fn truncate(s: &str, max: usize) -> String { fn truncate(s: &str, max: usize) -> String {
if s.len() <= max { s.to_string() } else { format!("{}", &s[..max]) } if s.len() <= max {
s.to_string()
} else {
format!("{}", &s[..max])
}
} }

View file

@ -14,6 +14,7 @@ pub struct FileRecord {
} }
impl FileRecord { impl FileRecord {
#[allow(dead_code)]
pub fn new(filename: impl Into<String>, path: impl Into<String>, description: Option<String>) -> Self { pub fn new(filename: impl Into<String>, path: impl Into<String>, description: Option<String>) -> Self {
Self { Self {
id: Uuid::new_v4().to_string(), id: Uuid::new_v4().to_string(),

View file

@ -1,5 +1,6 @@
use crate::gemini_client::{demo_text_embedding, DEMO_EMBED_DIM, generate_text}; use crate::gemini_client::{demo_text_embedding, generate_text_with_model, DEMO_EMBED_DIM};
use crate::models::{QueryRecord, QueryStatus}; use crate::models::{QueryRecord, QueryStatus};
use crate::vector;
use crate::vector_db::QdrantClient; use crate::vector_db::QdrantClient;
use anyhow::Result; use anyhow::Result;
use sqlx::MySqlPool; use sqlx::MySqlPool;
@ -84,15 +85,34 @@ impl Worker {
// Stage 2: embed query text // Stage 2: embed query text
let text = q.payload.get("q").and_then(|v| v.as_str()).unwrap_or(""); let text = q.payload.get("q").and_then(|v| v.as_str()).unwrap_or("");
let emb = demo_text_embedding(text).await?; let emb = demo_text_embedding(text).await?;
let top_k = q.payload.get("top_k").and_then(|v| v.as_u64()).unwrap_or(5) as usize; let top_k = q.payload.get("top_k").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
let top_k = top_k.max(1).min(20);
// Check cancellation // Check cancellation
if self.is_cancelled(&q.id).await? { return Ok(()); } if self.is_cancelled(&q.id).await? { return Ok(()); }
// Stage 3: search top-K in Qdrant // Stage 3: search top-K in Qdrant
let hits = self.qdrant.search_top_k(emb, top_k).await.unwrap_or_default(); let hits = match self.qdrant.search_top_k(emb.clone(), top_k).await {
let top_ids: Vec<String> = hits.iter().map(|(id, _)| id.clone()).collect(); Ok(list) if !list.is_empty() => list,
Ok(_) => Vec::new(),
Err(err) => {
error!("Qdrant search failed for query {}: {}", q.id, err);
Vec::new()
}
};
let hits = if hits.is_empty() {
match vector::query_top_k(&emb, top_k) {
Ok(fallback_ids) if !fallback_ids.is_empty() => {
info!("Using in-memory fallback for query {}", q.id);
fallback_ids.into_iter().map(|id| (id, 0.0)).collect()
}
_ => Vec::new(),
}
} else {
hits
};
// Check cancellation // Check cancellation
if self.is_cancelled(&q.id).await? { return Ok(()); } if self.is_cancelled(&q.id).await? { return Ok(()); }
@ -117,11 +137,23 @@ impl Worker {
// Stage 5: call Gemini to analyze relationships and propose follow-up details strictly from provided files // Stage 5: call Gemini to analyze relationships and propose follow-up details strictly from provided files
let relationships_prompt = build_relationships_prompt(text, &files_json); let relationships_prompt = build_relationships_prompt(text, &files_json);
let relationships = generate_text(&relationships_prompt).await.unwrap_or_else(|e| format!("[demo] relationships error: {}", e)); let (relationships, final_answer) = if files_json.is_empty() {
(
"No analyzed files are ready yet. Try seeding demo data or wait for processing to finish.".to_string(),
"I could not find any relevant documents yet. Once files finish analysis I will be able to answer.".to_string(),
)
} else {
let relationships = generate_text_with_model("gemini-2.5-pro", &relationships_prompt)
.await
.unwrap_or_else(|e| format!("[demo] relationships error: {}", e));
// Stage 6: final answer synthesis with strict constraints (no speculation; say unknown when insufficient) // Stage 6: final answer synthesis with strict constraints (no speculation; say unknown when insufficient)
let final_prompt = build_final_answer_prompt(text, &files_json, &relationships); let final_prompt = build_final_answer_prompt(text, &files_json, &relationships);
let final_answer = generate_text(&final_prompt).await.unwrap_or_else(|e| format!("[demo] final answer error: {}", e)); let final_answer = generate_text_with_model("gemini-2.5-pro", &final_prompt)
.await
.unwrap_or_else(|e| format!("[demo] final answer error: {}", e));
(relationships, final_answer)
};
// Stage 7: persist results // Stage 7: persist results
let result = serde_json::json!({ let result = serde_json::json!({