databao-context-engine 0.7.1.dev1__tar.gz → 0.7.1.dev2__tar.gz

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databao-context-engine
-Version: 0.7.1.dev1
+Version: 0.7.1.dev2
 Summary: Semantic context for your LLMs — generated automatically
 License-Expression: Apache-2.0 AND LicenseRef-Additional-Terms
 License-File: LICENSE.md

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "databao-context-engine"
-version = "0.7.1.dev1"
+version = "0.7.1.dev2"
 description = "Semantic context for your LLMs — generated automatically"
 readme = "README.md"
 requires-python = ">=3.11"

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/__init__.py

@@ -57,6 +57,7 @@ from databao_context_engine.plugins.databases.snowflake.config_file import (
     SnowflakeConfigFile,
     SnowflakeConnectionProperties,
     SnowflakeKeyPairAuth,
+    SnowflakeOAuthAuth,
     SnowflakePasswordAuth,
     SnowflakeSSOAuth,
 )
@@ -134,6 +135,7 @@ __all__ = [
     "SnowflakeConfigFile",
     "SnowflakeConnectionProperties",
     "SnowflakeSSOAuth",
+    "SnowflakeOAuthAuth",
     "SnowflakeKeyPairAuth",
     "SnowflakePasswordAuth",
     "SQLiteConfigFile",

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/build_sources/context_loader.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-from dataclasses import replace
 from typing import Any
 
 import yaml
-from pydantic import BaseModel, TypeAdapter
+from pydantic import TypeAdapter
 
 from databao_context_engine.build_sources.plugin_execution import BuiltDatasourceContext
 from databao_context_engine.datasources.datasource_context import (
@@ -45,14 +44,8 @@ def deserialize_built_context(
 ) -> BuiltDatasourceContext:
     """Parse a datasource output YAML payload and type the embedded context."""
     raw_context = yaml.safe_load(context.context)
-    built = TypeAdapter(BuiltDatasourceContext).validate_python(raw_context)
 
-    if isinstance(context_type, type) and issubclass(context_type, BaseModel):
-        typed_context: Any = context_type.model_validate(built.context)
-    else:
-        typed_context = TypeAdapter(context_type).validate_python(built.context)
-
-    return replace(built, context=typed_context)
+    return TypeAdapter(BuiltDatasourceContext[context_type]).validate_python(raw_context)  # type: ignore[valid-type]
 
 
 def _load_typed_built_context(

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/build_sources/plugin_execution.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, cast
+from typing import Any, Generic, TypeVar, cast
 
 from databao_context_engine.datasources.types import PreparedConfig, PreparedDatasource
 from databao_context_engine.pluginlib.build_plugin import (
@@ -10,9 +10,11 @@ from databao_context_engine.pluginlib.build_plugin import (
 from databao_context_engine.pluginlib.plugin_utils import execute_datasource_plugin, execute_file_plugin
 from databao_context_engine.project.layout import ProjectLayout
 
+ContextT = TypeVar("ContextT")
+
 
 @dataclass()
-class BuiltDatasourceContext:
+class BuiltDatasourceContext(Generic[ContextT]):
     """Dataclass defining the result of building a datasource's context."""
 
     datasource_id: str
@@ -25,7 +27,7 @@ class BuiltDatasourceContext:
     The type of the built data source
     """
 
-    context: Any
+    context: ContextT
     """
     A dictionary containing the actual context generated for the data source.
     This dictionary should be serializable in YAML format.

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/databao_context_engine.py

@@ -25,6 +25,7 @@ from databao_context_engine.plugins.databases.database_context_explorer import (
 from databao_context_engine.plugins.plugin_loader import DatabaoContextPluginLoader
 from databao_context_engine.project.layout import ProjectLayout, ensure_project_dir
 from databao_context_engine.search_context import search_context as search_context_internal
+from databao_context_engine.search_context.chunk_search_repository import ChunkType
 from databao_context_engine.search_context.search_service import ContextSearchMode
 
 
@@ -131,6 +132,7 @@ class DatabaoContextEngine:
         limit: int | None = None,
         datasource_ids: list[DatasourceId] | None = None,
         context_search_mode: ContextSearchMode | None = None,
+        chunk_types: list[ChunkType] | None = None,
     ) -> list[ContextSearchResult]:
         """Search in the available context for the closest matches to the given text.
 
@@ -139,6 +141,7 @@ class DatabaoContextEngine:
             limit: The maximum number of results to return. If None is provided, a default limit of 10 will be used.
             datasource_ids: If provided, the search results will only come from the datasources with these IDs.
             context_search_mode: Search strategy to use. Defaults to HYBRID_SEARCH if None is provided.
+            chunk_types: If provided, the search results will only come from the chunks of these types.
 
         Returns:
             A list of the results found for the search, sorted by score.
@@ -153,6 +156,7 @@ class DatabaoContextEngine:
             limit=limit,
             datasource_ids=datasource_ids,
             context_search_mode=context_search_mode,
+            chunk_types=chunk_types,
         )
 
         return [

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/pluginlib/build_plugin.py

@@ -12,10 +12,12 @@ class EmbeddableChunk:
     """A chunk that will be embedded as a vector and used when searching context from a given AI prompt.
 
     Attributes:
+        type: The type of the chunk, e.g. "table", "column" to allow for search by chunk type.
         embeddable_text: The text to embed as a vector for search usage
         content: The content to return as a response when the embedding has been selected in a search
     """
 
+    type: str | None = None
     embeddable_text: str
     keyword_indexable_text: str | None = None
     content: Any

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/plugins/databases/database_chunker.py

@@ -39,6 +39,7 @@ def build_database_chunks(result: DatabaseIntrospectionResult) -> list[Embeddabl
 
 def _create_table_chunk(catalog_name: str, schema_name: str, table: DatabaseTable) -> EmbeddableChunk:
     return EmbeddableChunk(
+        type="table",
         embeddable_text=_build_table_chunk_text(table),
         content=DatabaseTableChunkContent(
             catalog_name=catalog_name,
@@ -52,6 +53,7 @@ def _create_column_chunk(
     catalog_name: str, schema_name: str, table: DatabaseTable, column: DatabaseColumn
 ) -> EmbeddableChunk:
     return EmbeddableChunk(
+        type="column",
         embeddable_text=_build_column_chunk_text(table, column),
         content=DatabaseColumnChunkContent(
             catalog_name=catalog_name,

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/plugins/databases/snowflake/config_file.py

@@ -22,13 +22,17 @@ class SnowflakeSSOAuth(BaseModel):
     authenticator: str = Field(description='e.g. "externalbrowser"')
 
 
+class SnowflakeOAuthAuth(BaseModel):
+    token: Annotated[str, ConfigPropertyAnnotation(secret=True)]
+
+
 class SnowflakeConnectionProperties(BaseModel):
     account: Annotated[str, ConfigPropertyAnnotation(required=True)]
     warehouse: str | None = None
     database: str | None = None
     user: str | None = None
     role: str | None = None
-    auth: SnowflakePasswordAuth | SnowflakeKeyPairAuth | SnowflakeSSOAuth
+    auth: SnowflakePasswordAuth | SnowflakeKeyPairAuth | SnowflakeSSOAuth | SnowflakeOAuthAuth
     additional_properties: dict[str, Any] = {}
 
     def to_snowflake_kwargs(self) -> dict[str, Any]:

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/search_context/chunk_search_repository.py

@@ -1,6 +1,7 @@
 import logging
 from collections.abc import Sequence
 from dataclasses import dataclass
+from enum import Enum
 from typing import Any
 
 import duckdb
@@ -13,9 +14,17 @@ from databao_context_engine.pluginlib.build_plugin import DatasourceType
 logger = logging.getLogger(__name__)
 
 
+class ChunkType(str, Enum):
+    """Enum of the supported chunk types to search."""
+
+    TABLE = "table"
+    COLUMN = "column"
+
+
 @dataclass(kw_only=True, frozen=True)
 class VectorSearchCandidate:
     chunk_id: int
+    chunk_type: ChunkType | None
     display_text: str
     embeddable_text: str
     cosine_distance: float
@@ -26,6 +35,7 @@ class VectorSearchCandidate:
 @dataclass(kw_only=True, frozen=True)
 class Bm25SearchCandidate:
     chunk_id: int
+    chunk_type: ChunkType | None
     display_text: str
     embeddable_text: str
     bm25_score: float
@@ -65,6 +75,7 @@ class KeywordSearchScore:
 @dataclass(kw_only=True, frozen=True)
 class SearchResult:
     chunk_id: int
+    chunk_type: ChunkType | None
     display_text: str
     embeddable_text: str
     datasource_type: DatasourceType
@@ -90,6 +101,7 @@ class ChunkSearchRepository:
         dimension: int,
         limit: int,
         datasource_context_hashes: list[DatasourceContextHash],
+        chunk_types: list[ChunkType] | None = None,
     ) -> list[SearchResult]:
         """Read only similarity search on a specific embedding shard table."""
         vector_candidates = self._get_vector_candidates(
@@ -98,10 +110,12 @@ class ChunkSearchRepository:
             dimension=dimension,
             limit=limit,
             datasource_context_hashes=datasource_context_hashes,
+            chunk_types=chunk_types,
         )
         return [
             SearchResult(
                 chunk_id=candidate.chunk_id,
+                chunk_type=candidate.chunk_type,
                 display_text=candidate.display_text,
                 embeddable_text=candidate.embeddable_text,
                 datasource_type=candidate.datasource_type,
@@ -120,15 +134,26 @@ class ChunkSearchRepository:
         dimension: int,
         limit: int,
         datasource_context_hashes: list[DatasourceContextHash],
+        chunk_types: list[ChunkType] | None = None,
     ) -> list[VectorSearchCandidate]:
         """Read only vector candidates on a specific embedding shard table."""
         if not datasource_context_hashes:
             return []
 
         allowed_hashes_sql, hash_params = self._build_allowed_hashes_values(datasource_context_hashes)
+
+        chunk_types_param: list[list[ChunkType]]
+        if chunk_types:
+            search_candidates_chunk_type_filter = "WHERE c.chunk_type IN ?"
+            chunk_types_param = [chunk_types]
+        else:
+            chunk_types_param = []
+            search_candidates_chunk_type_filter = ""
+
         params: list[Any] = [
             *hash_params,
             list(search_vec),
+            *chunk_types_param,
             self._DEFAULT_DISTANCE_THRESHOLD,
             limit,
         ]
@@ -141,6 +166,7 @@ class ChunkSearchRepository:
             vector_candidates AS (
                 SELECT
                     c.chunk_id,
+                    c.chunk_type,
                     COALESCE(c.display_text, c.embeddable_text) AS display_text,
                     c.embeddable_text,
                     array_cosine_distance(e.vec, CAST(? AS FLOAT[{dimension}])) AS cosine_distance,
@@ -154,9 +180,11 @@ class ChunkSearchRepository:
                         ON h.datasource_id = ah.datasource_id
                         AND h.hash = ah.hash
                         AND h.hash_algorithm = ah.hash_algorithm
+                {search_candidates_chunk_type_filter}
             )
             SELECT
                 vc.chunk_id,
+                vc.chunk_type,
                 vc.display_text,
                 vc.embeddable_text,
                 vc.cosine_distance,
@@ -176,11 +204,12 @@ class ChunkSearchRepository:
         return [
             VectorSearchCandidate(
                 chunk_id=row[0],
-                display_text=row[1],
-                embeddable_text=row[2],
-                cosine_distance=row[3],
-                datasource_type=DatasourceType(full_type=row[4]),
-                datasource_id=DatasourceId.from_string_repr(row[5]),
+                chunk_type=ChunkType(row[1]) if row[1] else None,
+                display_text=row[2],
+                embeddable_text=row[3],
+                cosine_distance=row[4],
+                datasource_type=DatasourceType(full_type=row[5]),
+                datasource_id=DatasourceId.from_string_repr(row[6]),
             )
             for row in rows
         ]
@@ -195,6 +224,7 @@ class ChunkSearchRepository:
         dimension: int,
         limit: int,
         datasource_context_hashes: list[DatasourceContextHash],
+        chunk_types: list[ChunkType] | None = None,
     ) -> list[SearchResult]:
         """Hybrid retrieval combining vector similarity and BM25 with Reciprocal Rank Fusion (RRF).
 
@@ -208,12 +238,14 @@ class ChunkSearchRepository:
             dimension=dimension,
             limit=candidate_limit,
             datasource_context_hashes=datasource_context_hashes,
+            chunk_types=chunk_types,
         )
 
         bm25_candidates = self._get_bm25_candidates(
             query_text=search_text,
             limit=candidate_limit,
             datasource_context_hashes=datasource_context_hashes,
+            chunk_types=chunk_types,
         )
         return self._fuse_by_rrf(
             vector_candidates=vector_candidates,
@@ -228,17 +260,20 @@ class ChunkSearchRepository:
         query_text: str,
         limit: int,
         datasource_context_hashes: list[DatasourceContextHash],
+        chunk_types: list[ChunkType] | None = None,
     ) -> list[SearchResult]:
         """Read only BM25 search over chunk text."""
         bm25_candidates = self._get_bm25_candidates(
             query_text=query_text,
             limit=limit,
             datasource_context_hashes=datasource_context_hashes,
+            chunk_types=chunk_types,
         )
 
         return [
             SearchResult(
                 chunk_id=candidate.chunk_id,
+                chunk_type=candidate.chunk_type,
                 display_text=candidate.display_text,
                 embeddable_text=candidate.embeddable_text,
                 datasource_type=candidate.datasource_type,
@@ -255,12 +290,21 @@ class ChunkSearchRepository:
         query_text: str,
         limit: int,
         datasource_context_hashes: list[DatasourceContextHash],
+        chunk_types: list[ChunkType] | None = None,
     ) -> list[Bm25SearchCandidate]:
         if not datasource_context_hashes:
             return []
 
         allowed_hashes_sql, hash_params = self._build_allowed_hashes_values(datasource_context_hashes)
-        params: list[Any] = [*hash_params, query_text, limit]
+        chunk_types_param: list[list[ChunkType]]
+        if chunk_types:
+            search_candidates_chunk_type_filter = "WHERE c.chunk_type IN ?"
+            chunk_types_param = [chunk_types]
+        else:
+            chunk_types_param = []
+            search_candidates_chunk_type_filter = ""
+
+        params: list[Any] = [*hash_params, query_text, *chunk_types_param, limit]
 
         rows = self._conn.execute(
             f"""
@@ -270,6 +314,7 @@ class ChunkSearchRepository:
             bm25_candidates AS (
                 SELECT
                     c.chunk_id,
+                    c.chunk_type,
                     COALESCE(c.display_text, c.embeddable_text) AS display_text,
                     c.embeddable_text,
                     c.full_type,
@@ -285,9 +330,11 @@ class ChunkSearchRepository:
                         ON h.datasource_id = ah.datasource_id
                         AND h.hash = ah.hash
                         AND h.hash_algorithm = ah.hash_algorithm
+                {search_candidates_chunk_type_filter}
             )
             SELECT
                 b.chunk_id,
+                b.chunk_type,
                 b.display_text,
                 b.embeddable_text,
                 b.bm25_score,
@@ -307,11 +354,12 @@ class ChunkSearchRepository:
         return [
             Bm25SearchCandidate(
                 chunk_id=row[0],
-                display_text=row[1],
-                embeddable_text=row[2],
-                bm25_score=row[3],
-                datasource_type=DatasourceType(full_type=row[4]),
-                datasource_id=DatasourceId.from_string_repr(row[5]),
+                chunk_type=ChunkType(row[1]) if row[1] else None,
+                display_text=row[2],
+                embeddable_text=row[3],
+                bm25_score=row[4],
+                datasource_type=DatasourceType(full_type=row[5]),
+                datasource_id=DatasourceId.from_string_repr(row[6]),
             )
             for row in rows
         ]
@@ -365,6 +413,7 @@ class ChunkSearchRepository:
             results.append(
                 SearchResult(
                     chunk_id=chunk_id,
+                    chunk_type=data_candidate.chunk_type,
                     display_text=data_candidate.display_text,
                     embeddable_text=data_candidate.embeddable_text,
                     datasource_type=data_candidate.datasource_type,

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/search_context/search_runner.py

@@ -5,6 +5,7 @@ from databao_context_engine.datasources.datasource_context import (
 )
 from databao_context_engine.datasources.types import DatasourceId
 from databao_context_engine.project.layout import ProjectLayout
+from databao_context_engine.search_context.chunk_search_repository import ChunkType
 from databao_context_engine.search_context.search_service import RAG_MODE, ContextSearchMode, SearchContextService
 
 
@@ -18,6 +19,7 @@ def run_context_search(
     datasource_ids: list[DatasourceId] | None,
     rag_mode: RAG_MODE,
     context_search_mode: ContextSearchMode,
+    chunk_types: list[ChunkType] | None = None,
 ):
     context_hashes = (
         get_datasource_context_hashes(project_layout, datasource_ids)
@@ -33,4 +35,5 @@ def run_context_search(
         datasource_context_hashes=context_hashes,
         rag_mode=rag_mode,
         context_search_mode=context_search_mode,
+        chunk_types=chunk_types,
     )

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/search_context/search_service.py

@@ -8,6 +8,7 @@ from databao_context_engine.llm.embeddings.provider import EmbeddingProvider
 from databao_context_engine.llm.prompts.provider import PromptProvider
 from databao_context_engine.search_context.chunk_search_repository import (
     ChunkSearchRepository,
+    ChunkType,
     SearchResult,
 )
 from databao_context_engine.services.embedding_shard_resolver import EmbeddingShardResolver
@@ -55,6 +56,7 @@ class SearchContextService:
         limit: int | None = None,
         rag_mode: RAG_MODE,
         context_search_mode: ContextSearchMode,
+        chunk_types: list[ChunkType] | None = None,
     ) -> list[SearchResult]:
         if limit is None:
             limit = 10
@@ -65,6 +67,7 @@ class SearchContextService:
             limit=limit,
             rag_mode=rag_mode,
             context_search_mode=context_search_mode,
+            chunk_types=chunk_types,
         )
 
         logger.debug(f"Found {len(search_results)} search results")
@@ -90,6 +93,7 @@ class SearchContextService:
         limit: int,
         rag_mode: RAG_MODE,
         context_search_mode: ContextSearchMode,
+        chunk_types: list[ChunkType] | None = None,
     ) -> list[SearchResult]:
         if context_search_mode == ContextSearchMode.KEYWORD_SEARCH:
             query_text = self._rewrite_search_query(text) if rag_mode == RAG_MODE.REWRITE_QUERY else text
@@ -98,6 +102,7 @@ class SearchContextService:
                 query_text=query_text,
                 limit=limit,
                 datasource_context_hashes=datasource_context_hashes,
+                chunk_types=chunk_types,
             )
 
         table_name, dimension = self._shard_resolver.resolve(
@@ -130,6 +135,7 @@ class SearchContextService:
                     dimension=dimension,
                     limit=limit,
                     datasource_context_hashes=datasource_context_hashes,
+                    chunk_types=chunk_types,
                 )
             case ContextSearchMode.HYBRID_SEARCH:
                 return self._chunk_search_repo.search_chunks_with_hybrid_search(
@@ -139,6 +145,7 @@ class SearchContextService:
                     dimension=dimension,
                     limit=limit,
                     datasource_context_hashes=datasource_context_hashes,
+                    chunk_types=chunk_types,
                 )
 
     @perf.perf_span("search_context.rewrite_query")

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/search_context/search_wiring.py

@@ -14,7 +14,7 @@ from databao_context_engine.llm.factory import (
 from databao_context_engine.llm.prompts.provider import PromptProvider
 from databao_context_engine.plugins.plugin_loader import DatabaoContextPluginLoader
 from databao_context_engine.project.layout import ProjectLayout
-from databao_context_engine.search_context.chunk_search_repository import ChunkSearchRepository, SearchResult
+from databao_context_engine.search_context.chunk_search_repository import ChunkSearchRepository, ChunkType, SearchResult
 from databao_context_engine.search_context.search_runner import run_context_search
 from databao_context_engine.search_context.search_service import RAG_MODE, ContextSearchMode, SearchContextService
 from databao_context_engine.services.factories import create_shard_resolver
@@ -38,6 +38,7 @@ def search_context(
     limit: int | None,
     datasource_ids: list[DatasourceId] | None,
     context_search_mode: ContextSearchMode,
+    chunk_types: list[ChunkType] | None = None,
 ) -> list[SearchResult]:
     with open_duckdb_connection(project_layout.db_path) as conn:
         ollama_service = create_ollama_service()
@@ -62,6 +63,7 @@ def search_context(
             datasource_ids=datasource_ids,
             rag_mode=rag_mode,
             context_search_mode=context_search_mode,
+            chunk_types=chunk_types,
         )
 
 

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/services/persistence_service.py

@@ -135,7 +135,10 @@ class PersistenceService:
             full_type=full_type,
             datasource_id=datasource_id,
             datasource_context_hash_id=datasource_context_hash_id,
-            chunk_contents=[(ce.embedded_text, ce.display_text, ce.keyword_indexable_text) for ce in chunk_embeddings],
+            chunk_contents=[
+                (ce.embedded_text, ce.display_text, ce.keyword_indexable_text, ce.original_chunk.type)
+                for ce in chunk_embeddings
+            ],
         )
 
     @perf.perf_span("persistence.bulk_insert_embeddings")

databao_context_engine-0.7.1.dev2/src/databao_context_engine/storage/migrations/V04__add_chunk_type.sql

@@ -0,0 +1 @@
+ALTER TABLE chunk ADD COLUMN IF NOT EXISTS chunk_type TEXT;

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/storage/models.py

@@ -23,6 +23,7 @@ class ChunkDTO:
     display_text: Optional[str]
     created_at: datetime
     datasource_context_hash_id: int
+    chunk_type: str | None = None
 
 
 @dataclass(frozen=True)

{databao_context_engine-0.7.1.dev1 → databao_context_engine-0.7.1.dev2}/src/databao_context_engine/storage/repositories/chunk_repository.py

@@ -19,6 +19,7 @@ class ChunkRepository:
         self,
         *,
         full_type: str,
+        chunk_type: str | None = None,
         datasource_id: str,
         embeddable_text: str,
         display_text: Optional[str],
@@ -30,14 +31,15 @@ class ChunkRepository:
                 cur=self._conn,
                 sql="""
             INSERT INTO
-                chunk(full_type, datasource_id, embeddable_text, display_text, keyword_index_text, datasource_context_hash_id)
+                chunk(full_type, chunk_type, datasource_id, embeddable_text, display_text, keyword_index_text, datasource_context_hash_id)
             VALUES
-                (?, ?, ?, ?, ?, ?)
+                (?, ?, ?, ?, ?, ?, ?)
             RETURNING
                 *
             """,
                 params=[
                     full_type,
+                    chunk_type,
                     datasource_id,
                     embeddable_text,
                     display_text,
@@ -181,13 +183,13 @@ class ChunkRepository:
         *,
         full_type: str,
         datasource_id: str,
+        chunk_contents: Sequence[Tuple[str, Optional[str], str, Optional[str]]],
         datasource_context_hash_id: int,
-        chunk_contents: Sequence[Tuple[str, Optional[str], str]],
     ) -> Sequence[int]:
-        values_sql = ", ".join(["(?, ?, ?, ?, ?, ?)"] * len(chunk_contents))
+        values_sql = ", ".join(["(?, ?, ?, ?, ?, ?, ?)"] * len(chunk_contents))
         sql = f"""
             INSERT INTO
-                chunk(full_type, datasource_id, embeddable_text, display_text, keyword_index_text, datasource_context_hash_id)
+                chunk(full_type, chunk_type, datasource_id, embeddable_text, display_text, keyword_index_text, datasource_context_hash_id)
             VALUES
                 {values_sql}
             RETURNING
@@ -195,10 +197,11 @@ class ChunkRepository:
         """
 
         params: list[Any] = []
-        for embeddable_text, display_text, keyword_index_text in chunk_contents:
+        for embeddable_text, display_text, keyword_index_text, chunk_type in chunk_contents:
             params.extend(
                 [
                     full_type,
+                    chunk_type,
                     datasource_id,
                     embeddable_text,
                     display_text,
@@ -227,6 +230,7 @@ class ChunkRepository:
         return ChunkDTO(
             chunk_id=int(row["chunk_id"]),
             full_type=row["full_type"],
+            chunk_type=row["chunk_type"],
             datasource_id=row["datasource_id"],
             embeddable_text=row["embeddable_text"],
             display_text=row["display_text"],