Skip to content

Pipelines API

High-level orchestrators for AI workflows.

RAG Pipeline

maticlib.pipelines.rag_pipeline.RAGPipeline

RAGPipeline(
    llm_client,
    vector_index,
    prompt_name="rag_qa",
    use_hybrid=False,
    use_query_transform=False,
)

End-to-end RAG pipeline coordinating retrieval and generation.

Initializes the RAGPipeline.

Parameters:

Name Type Description Default
llm_client Any

Any LLM Client from maticlib.

required
vector_index BaseVectorIndex

A vector store index matching BaseVectorIndex.

required
prompt_name str

Name of prompt in the PromptRegistry. Default is 'rag_qa'.

'rag_qa'
use_hybrid bool

Whether to perform hybrid retrieval.

False
use_query_transform bool

Whether to expand the input question using the LLM.

False
Source code in maticlib/pipelines/rag_pipeline.py
def __init__(
    self,
    llm_client: Any,
    vector_index: BaseVectorIndex,
    prompt_name: str = "rag_qa",
    use_hybrid: bool = False,
    use_query_transform: bool = False,
):
    """
    Initializes the RAGPipeline.

    Args:
        llm_client: Any LLM Client from maticlib.
        vector_index: A vector store index matching BaseVectorIndex.
        prompt_name: Name of prompt in the PromptRegistry. Default is 'rag_qa'.
        use_hybrid: Whether to perform hybrid retrieval.
        use_query_transform: Whether to expand the input question using the LLM.
    """
    self.llm_client = llm_client
    self.vector_index = vector_index
    self.prompt_template = PromptRegistry.get(prompt_name)
    self.use_hybrid = use_hybrid
    self.use_query_transform = use_query_transform

    self.retriever = HybridRetriever(self.vector_index) if use_hybrid else None
    self.transformer = (
        QueryTransformer(self.llm_client) if use_query_transform else None
    )

generate

generate(question, k=4, keywords=None, trace=None)

Executes the RAG pipeline.

Parameters:

Name Type Description Default
question str

The natural language question to answer.

required
k int

Number of contexts to retrieve.

4
keywords Optional[List[str]]

Optional list of keywords for filtering.

None
trace Optional[PipelineTrace]

An optional trace container for observability.

None

Returns:

Type Description
str

The generated text answer string.

Source code in maticlib/pipelines/rag_pipeline.py
def generate(
    self,
    question: str,
    k: int = 4,
    keywords: Optional[List[str]] = None,
    trace: Optional[PipelineTrace] = None,
) -> str:
    """
    Executes the RAG pipeline.

    Args:
        question: The natural language question to answer.
        k: Number of contexts to retrieve.
        keywords: Optional list of keywords for filtering.
        trace: An optional trace container for observability.

    Returns:
        The generated text answer string.
    """
    # Step 1: Query Transformation (Optional)
    search_queries = [question]
    if self.transformer:
        transform_step = StepTrace(step_name="Query_Transform") if trace else None
        try:
            variations = self.transformer.generate_variations(question, n=2)
            search_queries.extend(variations)
        finally:
            if trace and transform_step:
                transform_step.end_time = time.time()
                trace.add_step(transform_step)

    # Step 2: Retrieval
    retrieve_step = StepTrace(step_name="Retrieval") if trace else None
    all_segments = []
    try:
        for q in set(search_queries):
            if self.retriever:
                segs = self.retriever.retrieve(q, k=k, keywords=keywords)
            else:
                segs = self.vector_index.similarity_search(q, k=k)
            all_segments.extend(segs)

        # Simple deduplication
        seen = set()
        unique_segments = []
        for s in all_segments:
            if s.segment_id not in seen:
                seen.add(s.segment_id)
                unique_segments.append(s)

        # Keep top k
        unique_segments = unique_segments[:k]
    finally:
        if trace and retrieve_step:
            retrieve_step.end_time = time.time()
            trace.add_step(retrieve_step)

    # Step 3: Formatting & Generation
    gen_step = StepTrace(step_name="Generation") if trace else None
    try:
        context = "\n\n".join(
            [f"--- Context ---\n{s.content}" for s in unique_segments]
        )
        prompt = self.prompt_template.format(context=context, question=question)

        res = self.llm_client.complete(prompt)

        # Extract token stats if available
        if hasattr(res, "prompt_tokens") and trace and gen_step:
            gen_step.prompt_tokens = getattr(res, "prompt_tokens", 0)
            gen_step.completion_tokens = getattr(res, "completion_tokens", 0)
            gen_step.total_tokens = getattr(res, "total_tokens", 0)

        return self.llm_client.get_text_response(res)
    finally:
        if trace and gen_step:
            gen_step.end_time = time.time()
            trace.add_step(gen_step)

Text2SQL Pipeline

maticlib.pipelines.text2sql_pipeline.Text2SQLPipeline

Text2SQLPipeline(
    llm_client,
    schema_loader,
    executor,
    connection_string,
    dialect="sqlite",
)

End-to-end Text2SQL pipeline.

Initializes the Text2SQLPipeline.

Parameters:

Name Type Description Default
llm_client Any

Any LLM client from maticlib.

required
schema_loader BaseSchemaLoader

A schema loader matching BaseSchemaLoader.

required
executor BaseExecutor

A DB executor matching BaseExecutor.

required
connection_string str

String URI for DB connection.

required
dialect str

The SQL dialect (e.g. 'sqlite', 'postgresql').

'sqlite'
Source code in maticlib/pipelines/text2sql_pipeline.py
def __init__(
    self,
    llm_client: Any,
    schema_loader: BaseSchemaLoader,
    executor: BaseExecutor,
    connection_string: str,
    dialect: str = "sqlite",
):
    """
    Initializes the Text2SQLPipeline.

    Args:
        llm_client: Any LLM client from maticlib.
        schema_loader: A schema loader matching BaseSchemaLoader.
        executor: A DB executor matching BaseExecutor.
        connection_string: String URI for DB connection.
        dialect: The SQL dialect (e.g. 'sqlite', 'postgresql').
    """
    self.llm_client = llm_client
    self.schema_loader = schema_loader
    self.executor = executor
    self.connection_string = connection_string
    self.dialect = dialect

    self.guard = SQLInjectionGuard(allowed_dialect=dialect)
    self.prompt_template = PromptRegistry.get("text2sql_generation")

execute

execute(question, trace=None)

Translates a question to SQL, validates it, and executes it. Returns columns and rows.

Parameters:

Name Type Description Default
question str

The natural language question to translate to SQL.

required
trace Optional[PipelineTrace]

Optional trace container for observability.

None

Returns:

Type Description
Tuple[List[str], List[Tuple]]

A tuple of column names (List[str]) and row tuples (List[Tuple]).

Source code in maticlib/pipelines/text2sql_pipeline.py
def execute(
    self, question: str, trace: Optional[PipelineTrace] = None
) -> Tuple[List[str], List[Tuple]]:
    """
    Translates a question to SQL, validates it, and executes it.
    Returns columns and rows.

    Args:
        question: The natural language question to translate to SQL.
        trace: Optional trace container for observability.

    Returns:
        A tuple of column names (List[str]) and row tuples (List[Tuple]).
    """
    # Step 1: Schema Loading
    schema_step = StepTrace(step_name="Schema_Load") if trace else None
    try:
        schema = self.schema_loader.load_schema(self.connection_string)
        schema_str = schema.to_prompt_string()
    finally:
        if trace and schema_step:
            schema_step.end_time = time.time()
            trace.add_step(schema_step)

    # Step 2: SQL Generation
    gen_step = StepTrace(step_name="SQL_Generation") if trace else None
    try:
        prompt = self.prompt_template.format(
            schema=schema_str, dialect=self.dialect, question=question
        )
        res = self.llm_client.complete(prompt)

        if hasattr(res, "prompt_tokens") and trace and gen_step:
            gen_step.prompt_tokens = getattr(res, "prompt_tokens", 0)
            gen_step.completion_tokens = getattr(res, "completion_tokens", 0)
            gen_step.total_tokens = getattr(res, "total_tokens", 0)

        raw_sql = self.llm_client.get_text_response(res)

        # Clean up markdown if LLM wrapped it
        raw_sql = raw_sql.strip("` \n")
        if raw_sql.startswith("sql"):
            raw_sql = raw_sql[3:].strip()

    finally:
        if trace and gen_step:
            gen_step.end_time = time.time()
            trace.add_step(gen_step)

    # Step 3: Validation and Execution
    exec_step = StepTrace(step_name="SQL_Execution") if trace else None
    try:
        safe_query = self.guard.validate_and_format(raw_sql)
        columns, rows = self.executor.execute(safe_query)
        return columns, rows
    except (SQLValidationError, SQLInjectionError) as e:
        if exec_step:
            exec_step.error = str(e)
        raise
    except Exception as e:
        if exec_step:
            exec_step.error = f"Execution failed: {e}"
        raise
    finally:
        if trace and exec_step:
            exec_step.end_time = time.time()
            trace.add_step(exec_step)