pgml/
pipeline.rs

1use anyhow::Context;
2use serde::Deserialize;
3use serde_json::json;
4use sqlx::{Executor, PgConnection, Pool, Postgres, Transaction};
5use std::collections::HashMap;
6use tracing::instrument;
7
8use crate::debug_sqlx_query;
9use crate::{
10    collection::ProjectInfo,
11    model::{Model, ModelRuntime},
12    models, queries, query_builder,
13    remote_embeddings::build_remote_embeddings,
14    splitter::Splitter,
15    types::{DateTime, Json, TryToNumeric},
16};
17
18#[cfg(feature = "rust_bridge")]
19use rust_bridge::{alias, alias_methods};
20
21#[cfg(feature = "python")]
22use crate::types::JsonPython;
23
24#[cfg(feature = "c")]
25use crate::languages::c::JsonC;
26
27type ParsedSchema = HashMap<String, FieldAction>;
28
29#[derive(Deserialize)]
30#[serde(deny_unknown_fields)]
31struct ValidSplitterAction {
32    model: Option<String>,
33    parameters: Option<Json>,
34}
35
36#[derive(Deserialize)]
37#[serde(deny_unknown_fields)]
38struct ValidEmbedAction {
39    model: String,
40    source: Option<String>,
41    parameters: Option<Json>,
42    hnsw: Option<Json>,
43}
44
45#[derive(Deserialize, Debug, Clone)]
46#[serde(deny_unknown_fields)]
47pub struct FullTextSearchAction {
48    configuration: String,
49}
50
51#[derive(Deserialize)]
52#[serde(deny_unknown_fields)]
53struct ValidFieldAction {
54    splitter: Option<ValidSplitterAction>,
55    semantic_search: Option<ValidEmbedAction>,
56    full_text_search: Option<FullTextSearchAction>,
57}
58
59#[allow(clippy::upper_case_acronyms)]
60#[derive(Debug, Clone)]
61pub struct HNSW {
62    m: u64,
63    ef_construction: u64,
64}
65
66impl Default for HNSW {
67    fn default() -> Self {
68        Self {
69            m: 16,
70            ef_construction: 64,
71        }
72    }
73}
74
75impl TryFrom<Json> for HNSW {
76    type Error = anyhow::Error;
77    fn try_from(value: Json) -> anyhow::Result<Self> {
78        let m = if !value["m"].is_null() {
79            value["m"]
80                .try_to_u64()
81                .context("hnsw.m must be an integer")?
82        } else {
83            16
84        };
85        let ef_construction = if !value["ef_construction"].is_null() {
86            value["ef_construction"]
87                .try_to_u64()
88                .context("hnsw.ef_construction must be an integer")?
89        } else {
90            64
91        };
92        Ok(Self { m, ef_construction })
93    }
94}
95
96#[derive(Debug, Clone)]
97pub struct SplitterAction {
98    pub model: Splitter,
99}
100
101#[derive(Debug, Clone)]
102pub struct SemanticSearchAction {
103    pub model: Model,
104    pub hnsw: HNSW,
105}
106
107#[derive(Debug, Clone)]
108pub struct FieldAction {
109    pub splitter: Option<SplitterAction>,
110    pub semantic_search: Option<SemanticSearchAction>,
111    pub full_text_search: Option<FullTextSearchAction>,
112}
113
114impl TryFrom<ValidFieldAction> for FieldAction {
115    type Error = anyhow::Error;
116    fn try_from(value: ValidFieldAction) -> Result<Self, Self::Error> {
117        let embed = value
118            .semantic_search
119            .map(|v| {
120                let model = Model::new(Some(v.model), v.source, v.parameters);
121                let hnsw = v
122                    .hnsw
123                    .map(HNSW::try_from)
124                    .unwrap_or_else(|| Ok(HNSW::default()))?;
125                anyhow::Ok(SemanticSearchAction { model, hnsw })
126            })
127            .transpose()?;
128        let splitter = value
129            .splitter
130            .map(|v| {
131                let splitter = Splitter::new(v.model, v.parameters);
132                anyhow::Ok(SplitterAction { model: splitter })
133            })
134            .transpose()?;
135        Ok(Self {
136            splitter,
137            semantic_search: embed,
138            full_text_search: value.full_text_search,
139        })
140    }
141}
142
143#[derive(Debug, Clone)]
144pub struct InvividualSyncStatus {
145    pub synced: i64,
146    pub not_synced: i64,
147    pub total: i64,
148}
149
150impl From<InvividualSyncStatus> for Json {
151    fn from(value: InvividualSyncStatus) -> Self {
152        serde_json::json!({
153            "synced": value.synced,
154            "not_synced": value.not_synced,
155            "total": value.total,
156        })
157        .into()
158    }
159}
160
161impl From<Json> for InvividualSyncStatus {
162    fn from(value: Json) -> Self {
163        Self {
164            synced: value["synced"]
165                .as_i64()
166                .expect("The synced field is not an integer"),
167            not_synced: value["not_synced"]
168                .as_i64()
169                .expect("The not_synced field is not an integer"),
170            total: value["total"]
171                .as_i64()
172                .expect("The total field is not an integer"),
173        }
174    }
175}
176
177#[derive(Debug, Clone)]
178#[allow(dead_code)]
179pub struct PipelineDatabaseData {
180    pub id: i64,
181    pub created_at: DateTime,
182}
183
184/// A pipeline that describes transformations to documents
185#[cfg_attr(feature = "rust_bridge", derive(alias))]
186#[derive(Debug, Clone)]
187pub struct Pipeline {
188    pub(crate) name: String,
189    pub(crate) schema: Option<Json>,
190    pub(crate) parsed_schema: Option<ParsedSchema>,
191    database_data: Option<PipelineDatabaseData>,
192}
193
194fn json_to_schema(schema: &Json) -> anyhow::Result<ParsedSchema> {
195    schema
196        .as_object()
197        .context("Schema object must be a JSON object")?
198        .iter()
199        .try_fold(ParsedSchema::new(), |mut acc, (key, value)| {
200            if acc.contains_key(key) {
201                Err(anyhow::anyhow!("Schema contains duplicate keys"))
202            } else {
203                // First lets deserialize it normally
204                let action: ValidFieldAction = serde_json::from_value(value.to_owned())?;
205                // Now lets actually build the models and splitters
206                acc.insert(key.to_owned(), action.try_into()?);
207                Ok(acc)
208            }
209        })
210}
211
212#[cfg_attr(feature = "rust_bridge", alias_methods(new))]
213impl Pipeline {
214    /// Creates a [Pipeline]
215    ///
216    /// # Arguments
217    /// * `name` - The name of the pipeline
218    /// * `schema` - The schema of the pipeline. This is a JSON object where the keys are the field names and the values are the field actions.
219    pub fn new(name: &str, schema: Option<Json>) -> anyhow::Result<Self> {
220        let parsed_schema = schema.as_ref().map(json_to_schema).transpose()?;
221        Ok(Self {
222            name: name.to_string(),
223            schema,
224            parsed_schema,
225            database_data: None,
226        })
227    }
228
229    /// Gets the status of the [Pipeline]
230    #[instrument(skip(self))]
231    pub(crate) async fn get_status(
232        &mut self,
233        project_info: &ProjectInfo,
234        pool: &Pool<Postgres>,
235    ) -> anyhow::Result<Json> {
236        let parsed_schema = self
237            .parsed_schema
238            .as_ref()
239            .context("Pipeline must have schema to get status")?;
240
241        let mut results = json!({});
242
243        let schema = format!("{}_{}", project_info.name, self.name);
244        let documents_table_name = format!("{}.documents", project_info.name);
245        for (key, value) in parsed_schema.iter() {
246            let chunks_table_name = format!("{schema}.{key}_chunks");
247
248            results[key] = json!({});
249
250            if value.splitter.is_some() {
251                let chunks_status: (Option<i64>, Option<i64>) = sqlx::query_as(&query_builder!(
252                    "SELECT (SELECT COUNT(DISTINCT document_id) FROM %s), COUNT(id) FROM %s",
253                    chunks_table_name,
254                    documents_table_name
255                ))
256                .fetch_one(pool)
257                .await?;
258                results[key]["chunks"] = json!({
259                    "synced": chunks_status.0.unwrap_or(0),
260                    "not_synced": chunks_status.1.unwrap_or(0) - chunks_status.0.unwrap_or(0),
261                    "total": chunks_status.1.unwrap_or(0),
262                });
263            }
264
265            if value.semantic_search.is_some() {
266                let embeddings_table_name = format!("{schema}.{key}_embeddings");
267                let embeddings_status: (Option<i64>, Option<i64>) =
268                    sqlx::query_as(&query_builder!(
269                        "SELECT (SELECT count(*) FROM %s), (SELECT count(*) FROM %s)",
270                        embeddings_table_name,
271                        chunks_table_name
272                    ))
273                    .fetch_one(pool)
274                    .await?;
275                results[key]["embeddings"] = json!({
276                    "synced": embeddings_status.0.unwrap_or(0),
277                    "not_synced": embeddings_status.1.unwrap_or(0) - embeddings_status.0.unwrap_or(0),
278                    "total": embeddings_status.1.unwrap_or(0),
279                });
280            }
281
282            if value.full_text_search.is_some() {
283                let tsvectors_table_name = format!("{schema}.{key}_tsvectors");
284                let tsvectors_status: (Option<i64>, Option<i64>) = sqlx::query_as(&query_builder!(
285                    "SELECT (SELECT count(*) FROM %s), (SELECT count(*) FROM %s)",
286                    tsvectors_table_name,
287                    chunks_table_name
288                ))
289                .fetch_one(pool)
290                .await?;
291                results[key]["tsvectors"] = json!({
292                    "synced": tsvectors_status.0.unwrap_or(0),
293                    "not_synced": tsvectors_status.1.unwrap_or(0) - tsvectors_status.0.unwrap_or(0),
294                    "total": tsvectors_status.1.unwrap_or(0),
295                });
296            }
297        }
298        Ok(results.into())
299    }
300
301    #[instrument(skip(self))]
302    pub(crate) async fn verify_in_database(
303        &mut self,
304        project_info: &ProjectInfo,
305        throw_if_exists: bool,
306        pool: &Pool<Postgres>,
307    ) -> anyhow::Result<()> {
308        if self.database_data.is_none() {
309            let pipeline: Option<models::Pipeline> = sqlx::query_as(&query_builder!(
310                "SELECT * FROM %s WHERE name = $1",
311                format!("{}.pipelines", project_info.name)
312            ))
313            .bind(&self.name)
314            .fetch_optional(pool)
315            .await?;
316
317            let pipeline = if let Some(pipeline) = pipeline {
318                if throw_if_exists {
319                    anyhow::bail!("Pipeline {} already exists. You do not need to add this pipeline to the collection as it has already been added.", pipeline.name);
320                }
321
322                let mut parsed_schema = json_to_schema(&pipeline.schema)?;
323
324                for (_key, value) in parsed_schema.iter_mut() {
325                    if let Some(splitter) = &mut value.splitter {
326                        splitter
327                            .model
328                            .verify_in_database(project_info, false, pool)
329                            .await?;
330                    }
331                    if let Some(embed) = &mut value.semantic_search {
332                        embed
333                            .model
334                            .verify_in_database(project_info, false, pool)
335                            .await?;
336                    }
337                }
338                self.schema = Some(pipeline.schema.clone());
339                self.parsed_schema = Some(parsed_schema);
340
341                pipeline
342            } else {
343                let schema = self
344                    .schema
345                    .as_ref()
346                    .context("Pipeline must have schema to store in database")?;
347                let mut parsed_schema = json_to_schema(schema)?;
348
349                for (_key, value) in parsed_schema.iter_mut() {
350                    if let Some(splitter) = &mut value.splitter {
351                        splitter
352                            .model
353                            .verify_in_database(project_info, false, pool)
354                            .await?;
355                    }
356                    if let Some(embed) = &mut value.semantic_search {
357                        embed
358                            .model
359                            .verify_in_database(project_info, false, pool)
360                            .await?;
361                    }
362                }
363                self.parsed_schema = Some(parsed_schema);
364
365                // Here we actually insert the pipeline into the collection.pipelines table
366                // and create the collection_pipeline schema and required tables
367                let mut transaction = pool.begin().await?;
368                let pipeline = sqlx::query_as(&query_builder!(
369                    "INSERT INTO %s (name, schema) VALUES ($1, $2) RETURNING *",
370                    format!("{}.pipelines", project_info.name)
371                ))
372                .bind(&self.name)
373                .bind(&self.schema)
374                .fetch_one(&mut *transaction)
375                .await?;
376                self.create_tables(project_info, &mut transaction).await?;
377                transaction.commit().await?;
378
379                pipeline
380            };
381            self.database_data = Some(PipelineDatabaseData {
382                id: pipeline.id,
383                created_at: pipeline.created_at,
384            })
385        }
386        Ok(())
387    }
388
389    #[instrument(skip(self))]
390    async fn create_tables(
391        &mut self,
392        project_info: &ProjectInfo,
393        transaction: &mut Transaction<'_, Postgres>,
394    ) -> anyhow::Result<()> {
395        let collection_name = &project_info.name;
396        let documents_table_name = format!("{}.documents", collection_name);
397
398        let schema = format!("{}_{}", collection_name, self.name);
399
400        transaction
401            .execute(query_builder!("CREATE SCHEMA IF NOT EXISTS %s", schema).as_str())
402            .await?;
403
404        let parsed_schema = self
405            .parsed_schema
406            .as_ref()
407            .context("Pipeline must have schema to create_tables")?;
408
409        let searches_table_name = format!("{schema}.searches");
410        transaction
411            .execute(
412                query_builder!(
413                    queries::CREATE_PIPELINES_SEARCHES_TABLE,
414                    searches_table_name
415                )
416                .as_str(),
417            )
418            .await?;
419
420        let search_results_table_name = format!("{schema}.search_results");
421        transaction
422            .execute(
423                query_builder!(
424                    queries::CREATE_PIPELINES_SEARCH_RESULTS_TABLE,
425                    search_results_table_name,
426                    &searches_table_name,
427                    &documents_table_name
428                )
429                .as_str(),
430            )
431            .await?;
432        transaction
433            .execute(
434                query_builder!(
435                    queries::CREATE_INDEX,
436                    "",
437                    "search_results_search_id_rank_index",
438                    search_results_table_name,
439                    "search_id, rank"
440                )
441                .as_str(),
442            )
443            .await?;
444
445        let search_events_table_name = format!("{schema}.search_events");
446        transaction
447            .execute(
448                query_builder!(
449                    queries::CREATE_PIPELINES_SEARCH_EVENTS_TABLE,
450                    search_events_table_name,
451                    &search_results_table_name
452                )
453                .as_str(),
454            )
455            .await?;
456
457        for (key, value) in parsed_schema.iter() {
458            let chunks_table_name = format!("{}.{}_chunks", schema, key);
459            transaction
460                .execute(
461                    query_builder!(
462                        queries::CREATE_CHUNKS_TABLE,
463                        chunks_table_name,
464                        documents_table_name
465                    )
466                    .as_str(),
467                )
468                .await?;
469            let index_name = format!("{}_pipeline_chunk_document_id_index", key);
470            transaction
471                .execute(
472                    query_builder!(
473                        queries::CREATE_INDEX,
474                        "",
475                        index_name,
476                        chunks_table_name,
477                        "document_id"
478                    )
479                    .as_str(),
480                )
481                .await?;
482
483            if let Some(embed) = &value.semantic_search {
484                let embeddings_table_name = format!("{}.{}_embeddings", schema, key);
485                let embedding_length = match &embed.model.runtime {
486                    ModelRuntime::Python => {
487                        let embedding: (Vec<f32>,) = sqlx::query_as(
488                                    "SELECT embedding from pgml.embed(transformer => $1, text => 'Hello, World!', kwargs => $2) as embedding")
489                                    .bind(&embed.model.name)
490                                    .bind(&embed.model.parameters)
491                                    .fetch_one(&mut **transaction).await?;
492                        embedding.0.len() as i64
493                    }
494                    t => {
495                        let remote_embeddings = build_remote_embeddings(
496                            t.to_owned(),
497                            &embed.model.name,
498                            Some(&embed.model.parameters),
499                        )?;
500                        remote_embeddings.get_embedding_size().await?
501                    }
502                };
503
504                // Create the embeddings table
505                sqlx::query(&query_builder!(
506                    queries::CREATE_EMBEDDINGS_TABLE,
507                    &embeddings_table_name,
508                    chunks_table_name,
509                    embedding_length
510                ))
511                .execute(&mut **transaction)
512                .await?;
513                let index_name = format!("{}_pipeline_embedding_chunk_id_index", key);
514                transaction
515                    .execute(
516                        query_builder!(
517                            queries::CREATE_INDEX,
518                            "",
519                            index_name,
520                            &embeddings_table_name,
521                            "chunk_id"
522                        )
523                        .as_str(),
524                    )
525                    .await?;
526                let index_with_parameters = format!(
527                    "WITH (m = {}, ef_construction = {})",
528                    embed.hnsw.m, embed.hnsw.ef_construction
529                );
530                let index_name = format!("{}_pipeline_embedding_hnsw_vector_index", key);
531                transaction
532                    .execute(
533                        query_builder!(
534                            queries::CREATE_INDEX_USING_HNSW,
535                            "",
536                            index_name,
537                            &embeddings_table_name,
538                            "embedding vector_cosine_ops",
539                            index_with_parameters
540                        )
541                        .as_str(),
542                    )
543                    .await?;
544            }
545
546            // Create the tsvectors table
547            if value.full_text_search.is_some() {
548                let tsvectors_table_name = format!("{}.{}_tsvectors", schema, key);
549                transaction
550                    .execute(
551                        query_builder!(
552                            queries::CREATE_CHUNKS_TSVECTORS_TABLE,
553                            tsvectors_table_name,
554                            chunks_table_name
555                        )
556                        .as_str(),
557                    )
558                    .await?;
559                let index_name = format!("{}_pipeline_tsvector_chunk_id_index", key);
560                transaction
561                    .execute(
562                        query_builder!(
563                            queries::CREATE_INDEX,
564                            "",
565                            index_name,
566                            tsvectors_table_name,
567                            "chunk_id"
568                        )
569                        .as_str(),
570                    )
571                    .await?;
572                let index_name = format!("{}_pipeline_tsvector_index", key);
573                transaction
574                    .execute(
575                        query_builder!(
576                            queries::CREATE_INDEX_USING_GIN,
577                            "",
578                            index_name,
579                            tsvectors_table_name,
580                            "ts"
581                        )
582                        .as_str(),
583                    )
584                    .await?;
585            }
586        }
587        Ok(())
588    }
589
590    #[instrument(skip(self))]
591    pub(crate) async fn sync_documents(
592        &mut self,
593        document_ids: Vec<i64>,
594        project_info: &ProjectInfo,
595        transaction: &mut Transaction<'static, Postgres>,
596    ) -> anyhow::Result<()> {
597        // We are assuming we have manually verified the pipeline before doing this
598        let parsed_schema = self
599            .parsed_schema
600            .as_ref()
601            .context("Pipeline must have schema to execute")?;
602
603        for (key, value) in parsed_schema.iter() {
604            let chunk_ids = self
605                .sync_chunks_for_documents(
606                    key,
607                    value.splitter.as_ref().map(|v| &v.model),
608                    &document_ids,
609                    project_info,
610                    transaction,
611                )
612                .await?;
613            if !chunk_ids.is_empty() {
614                if let Some(embed) = &value.semantic_search {
615                    self.sync_embeddings_for_chunks(
616                        key,
617                        &embed.model,
618                        &chunk_ids,
619                        project_info,
620                        transaction,
621                    )
622                    .await?;
623                }
624                if let Some(full_text_search) = &value.full_text_search {
625                    self.sync_tsvectors_for_chunks(
626                        key,
627                        &full_text_search.configuration,
628                        &chunk_ids,
629                        project_info,
630                        transaction,
631                    )
632                    .await?;
633                }
634            }
635        }
636        Ok(())
637    }
638
639    #[instrument(skip(self))]
640    async fn sync_chunks_for_documents(
641        &self,
642        key: &str,
643        splitter: Option<&Splitter>,
644        document_ids: &Vec<i64>,
645        project_info: &ProjectInfo,
646        transaction: &mut Transaction<'static, Postgres>,
647    ) -> anyhow::Result<Vec<i64>> {
648        let chunks_table_name = format!("{}_{}.{}_chunks", project_info.name, self.name, key);
649        let documents_table_name = format!("{}.documents", project_info.name);
650        let json_key_query = format!("document->>'{}'", key);
651
652        if let Some(splitter) = splitter {
653            let splitter_database_data = splitter
654                .database_data
655                .as_ref()
656                .context("Splitter must be verified to sync chunks")?;
657            let query = query_builder!(
658                queries::GENERATE_CHUNKS_FOR_DOCUMENT_IDS_WITH_SPLITTER,
659                &json_key_query,
660                documents_table_name,
661                &chunks_table_name,
662                &chunks_table_name,
663                &chunks_table_name
664            );
665            debug_sqlx_query!(
666                GENERATE_CHUNKS_FOR_DOCUMENT_IDS_WITH_SPLITTER,
667                query,
668                splitter_database_data.id,
669                document_ids
670            );
671            sqlx::query_scalar(&query)
672                .bind(splitter_database_data.id)
673                .bind(document_ids)
674                .fetch_all(&mut **transaction)
675                .await
676                .map_err(anyhow::Error::msg)
677        } else {
678            let query = query_builder!(
679                queries::GENERATE_CHUNKS_FOR_DOCUMENT_IDS,
680                &chunks_table_name,
681                &json_key_query,
682                &documents_table_name,
683                &chunks_table_name,
684                &json_key_query
685            );
686            debug_sqlx_query!(GENERATE_CHUNKS_FOR_DOCUMENT_IDS, query, document_ids);
687            sqlx::query_scalar(&query)
688                .bind(document_ids)
689                .fetch_all(&mut **transaction)
690                .await
691                .map_err(anyhow::Error::msg)
692        }
693    }
694
695    #[instrument(skip(self))]
696    async fn sync_embeddings_for_chunks(
697        &self,
698        key: &str,
699        model: &Model,
700        chunk_ids: &Vec<i64>,
701        project_info: &ProjectInfo,
702        transaction: &mut Transaction<'static, Postgres>,
703    ) -> anyhow::Result<()> {
704        // Remove the stored name from the parameters
705        let mut parameters = model.parameters.clone();
706        parameters
707            .as_object_mut()
708            .context("Model parameters must be an object")?
709            .remove("name");
710
711        let chunks_table_name = format!("{}_{}.{}_chunks", project_info.name, self.name, key);
712        let embeddings_table_name =
713            format!("{}_{}.{}_embeddings", project_info.name, self.name, key);
714
715        match model.runtime {
716            ModelRuntime::Python => {
717                let query = query_builder!(
718                    queries::GENERATE_EMBEDDINGS_FOR_CHUNK_IDS,
719                    embeddings_table_name,
720                    chunks_table_name
721                );
722                debug_sqlx_query!(
723                    GENERATE_EMBEDDINGS_FOR_CHUNK_IDS,
724                    query,
725                    model.name,
726                    parameters.0,
727                    chunk_ids
728                );
729                sqlx::query(&query)
730                    .bind(&model.name)
731                    .bind(&parameters)
732                    .bind(chunk_ids)
733                    .execute(&mut **transaction)
734                    .await?;
735            }
736            r => {
737                let remote_embeddings = build_remote_embeddings(r, &model.name, Some(&parameters))?;
738                remote_embeddings
739                    .generate_embeddings(
740                        &embeddings_table_name,
741                        &chunks_table_name,
742                        Some(chunk_ids),
743                        transaction,
744                    )
745                    .await?;
746            }
747        }
748        Ok(())
749    }
750
751    #[instrument(skip(self))]
752    async fn sync_tsvectors_for_chunks(
753        &self,
754        key: &str,
755        configuration: &str,
756        chunk_ids: &Vec<i64>,
757        project_info: &ProjectInfo,
758        transaction: &mut Transaction<'static, Postgres>,
759    ) -> anyhow::Result<()> {
760        let chunks_table_name = format!("{}_{}.{}_chunks", project_info.name, self.name, key);
761        let tsvectors_table_name = format!("{}_{}.{}_tsvectors", project_info.name, self.name, key);
762        let query = query_builder!(
763            queries::GENERATE_TSVECTORS_FOR_CHUNK_IDS,
764            tsvectors_table_name,
765            configuration,
766            chunks_table_name
767        );
768        debug_sqlx_query!(GENERATE_TSVECTORS_FOR_CHUNK_IDS, query, chunk_ids);
769        sqlx::query(&query)
770            .bind(chunk_ids)
771            .execute(&mut **transaction)
772            .await?;
773        Ok(())
774    }
775
776    #[instrument(skip(self))]
777    pub(crate) async fn resync(
778        &mut self,
779        project_info: &ProjectInfo,
780        connection: &mut PgConnection,
781    ) -> anyhow::Result<()> {
782        // We are assuming we have manually verified the pipeline before doing this
783        let parsed_schema = self
784            .parsed_schema
785            .as_ref()
786            .context("Pipeline must have schema to execute")?;
787        // Before doing any syncing, delete all old and potentially outdated documents
788        for (key, _value) in parsed_schema.iter() {
789            let chunks_table_name = format!("{}_{}.{}_chunks", project_info.name, self.name, key);
790            connection
791                .execute(query_builder!("DELETE FROM %s CASCADE", chunks_table_name).as_str())
792                .await?;
793        }
794        for (key, value) in parsed_schema.iter() {
795            self.resync_chunks(
796                key,
797                value.splitter.as_ref().map(|v| &v.model),
798                project_info,
799                connection,
800            )
801            .await?;
802            if let Some(embed) = &value.semantic_search {
803                self.resync_embeddings(key, &embed.model, project_info, connection)
804                    .await?;
805            }
806            if let Some(full_text_search) = &value.full_text_search {
807                self.resync_tsvectors(
808                    key,
809                    &full_text_search.configuration,
810                    project_info,
811                    connection,
812                )
813                .await?;
814            }
815        }
816        Ok(())
817    }
818
819    #[instrument(skip(self))]
820    async fn resync_chunks(
821        &self,
822        key: &str,
823        splitter: Option<&Splitter>,
824        project_info: &ProjectInfo,
825        connection: &mut PgConnection,
826    ) -> anyhow::Result<()> {
827        let chunks_table_name = format!("{}_{}.{}_chunks", project_info.name, self.name, key);
828        let documents_table_name = format!("{}.documents", project_info.name);
829        let json_key_query = format!("document->>'{}'", key);
830
831        if let Some(splitter) = splitter {
832            let splitter_database_data = splitter
833                .database_data
834                .as_ref()
835                .context("Splitter must be verified to sync chunks")?;
836            let query = query_builder!(
837                queries::GENERATE_CHUNKS_WITH_SPLITTER,
838                &json_key_query,
839                &documents_table_name,
840                &chunks_table_name,
841                &chunks_table_name
842            );
843            debug_sqlx_query!(
844                GENERATE_CHUNKS_WITH_SPLITTER,
845                query,
846                splitter_database_data.id
847            );
848            sqlx::query(&query)
849                .bind(splitter_database_data.id)
850                .execute(connection)
851                .await?;
852        } else {
853            let query = query_builder!(
854                queries::GENERATE_CHUNKS,
855                &chunks_table_name,
856                &json_key_query,
857                &documents_table_name
858            );
859            debug_sqlx_query!(GENERATE_CHUNKS, query);
860            sqlx::query(&query).execute(connection).await?;
861        }
862        Ok(())
863    }
864
865    #[instrument(skip(self))]
866    async fn resync_embeddings(
867        &self,
868        key: &str,
869        model: &Model,
870        project_info: &ProjectInfo,
871        connection: &mut PgConnection,
872    ) -> anyhow::Result<()> {
873        // Remove the stored name from the parameters
874        let mut parameters = model.parameters.clone();
875        parameters
876            .as_object_mut()
877            .context("Model parameters must be an object")?
878            .remove("name");
879
880        let chunks_table_name = format!("{}_{}.{}_chunks", project_info.name, self.name, key);
881        let embeddings_table_name =
882            format!("{}_{}.{}_embeddings", project_info.name, self.name, key);
883
884        match model.runtime {
885            ModelRuntime::Python => {
886                let query = query_builder!(
887                    queries::GENERATE_EMBEDDINGS,
888                    embeddings_table_name,
889                    chunks_table_name
890                );
891                debug_sqlx_query!(GENERATE_EMBEDDINGS, query, model.name, parameters.0);
892                sqlx::query(&query)
893                    .bind(&model.name)
894                    .bind(&parameters)
895                    .execute(connection)
896                    .await?;
897            }
898            r => {
899                let remote_embeddings = build_remote_embeddings(r, &model.name, Some(&parameters))?;
900                remote_embeddings
901                    .generate_embeddings(
902                        &embeddings_table_name,
903                        &chunks_table_name,
904                        None,
905                        connection,
906                    )
907                    .await?;
908            }
909        }
910        Ok(())
911    }
912
913    #[instrument(skip(self))]
914    async fn resync_tsvectors(
915        &self,
916        key: &str,
917        configuration: &str,
918        project_info: &ProjectInfo,
919        connection: &mut PgConnection,
920    ) -> anyhow::Result<()> {
921        let chunks_table_name = format!("{}_{}.{}_chunks", project_info.name, self.name, key);
922        let tsvectors_table_name = format!("{}_{}.{}_tsvectors", project_info.name, self.name, key);
923
924        let query = query_builder!(
925            queries::GENERATE_TSVECTORS,
926            tsvectors_table_name,
927            configuration,
928            chunks_table_name
929        );
930        debug_sqlx_query!(GENERATE_TSVECTORS, query);
931        sqlx::query(&query).execute(connection).await?;
932        Ok(())
933    }
934
935    #[instrument(skip(self))]
936    pub(crate) async fn get_parsed_schema(
937        &mut self,
938        project_info: &ProjectInfo,
939        pool: &Pool<Postgres>,
940    ) -> anyhow::Result<ParsedSchema> {
941        self.verify_in_database(project_info, false, pool).await?;
942        Ok(self.parsed_schema.as_ref().unwrap().clone())
943    }
944
945    #[instrument]
946    pub(crate) async fn create_pipelines_table(
947        project_info: &ProjectInfo,
948        conn: &mut PgConnection,
949    ) -> anyhow::Result<()> {
950        let pipelines_table_name = format!("{}.pipelines", project_info.name);
951        sqlx::query(&query_builder!(
952            queries::PIPELINES_TABLE,
953            pipelines_table_name
954        ))
955        .execute(&mut *conn)
956        .await?;
957        conn.execute(
958            query_builder!(
959                queries::CREATE_INDEX,
960                "",
961                "pipeline_name_index",
962                pipelines_table_name,
963                "name"
964            )
965            .as_str(),
966        )
967        .await?;
968        Ok(())
969    }
970}
971
972impl TryFrom<models::Pipeline> for Pipeline {
973    type Error = anyhow::Error;
974    fn try_from(value: models::Pipeline) -> anyhow::Result<Self> {
975        let parsed_schema = json_to_schema(&value.schema).unwrap();
976        // NOTE: We do not set the database data here even though we have it
977        // self.verify_in_database() also verifies all models in the schema so we don't want to set it here
978        Ok(Self {
979            name: value.name,
980            schema: Some(value.schema),
981            parsed_schema: Some(parsed_schema),
982            database_data: None,
983        })
984    }
985}
pgml/pipeline.rs

pgml/
pipeline.rs
