feat(cli): complete sync pipeline

Author: JohnTrunix

.gitignore

@@ -7,6 +7,7 @@ __pycache__
 dist/
 
 .pytest_cache
+.mypy_cache
 .ruff_cache
 .venv
 *.env

grimoire.yaml

@@ -16,4 +16,10 @@ llm:
   code_chunk_overlap: 128
 include_project: true
 project_src: grimoire
-sources: null
+sources:
+  - url: https://github.com/fastapi/typer
+    include_md: true
+    include_code: false
+  - url: https://github.com/pydantic/pydantic
+    include_md: true
+    include_code: false

grimoire/configuration.py

@@ -85,7 +85,7 @@ class Source(BaseModel):
     """
 
     url: str
-    branch: str = "default"
+    branch: str | None = None
     include_md: bool = True
     include_code: bool = False
 

grimoire/init.py

@@ -7,6 +7,7 @@
     DBConfiguration,
     LLMConfiguration,
     ProjectConfiguration,
+    Source,
 )
 from grimoire.helpers.typer import green_text
 
@@ -71,12 +72,20 @@ def init(
         raise typer.Exit(code=0)
 
     config: ProjectConfiguration = get_project_config(path)
+    config.sources = [
+        Source(
+            url="https://github.com/BDP25/grimoire",
+            branch="main",
+            include_md=True,
+            include_code=False,
+        )
+    ]
     config.save_to_yaml(file_path)
 
     success_message = f"""
     {green_text("Grimoire project initialized successfully! 🎉")}
 
-    Run `grimoire update` to update all your projects dependencies.
+    Add your sources to the project by modifying the {green_text("sources")} field in the configuration file.
 
     Configuration: {file_path}
     Note: directly modify the "{CONFIG_FILE_NAME}" configuration to your needs.

grimoire/main.py

@@ -5,7 +5,6 @@
 from grimoire.flush import flush_cli
 from grimoire.init import init_cli
 from grimoire.sync import sync_cli
-from grimoire.update import update_cli
 from grimoire.verify import verify_cli
 
 CLI_NAME = "grimoire"
@@ -27,7 +26,6 @@
 cli.add_typer(flush_cli, name=None)
 cli.add_typer(init_cli, name=None)
 cli.add_typer(sync_cli, name=None)
-cli.add_typer(update_cli, name=None)
 cli.add_typer(verify_cli, name=None)
 
 

grimoire/sync.py

@@ -5,6 +5,7 @@
 import psycopg
 import typer
 from git import Repo
+from langchain_core.documents import Document
 from langchain_postgres import PGVector
 from rich.progress import track
 
@@ -14,6 +15,7 @@
     get_recursive_config,
 )
 from grimoire.helpers.ingestion import code_ingestion, text_ingestion
+from grimoire.helpers.typer import green_text
 from grimoire.helpers.vectorstore import (
     clear_collection,
     setup_vectorstore,
@@ -24,6 +26,25 @@
 sync_cli = typer.Typer()
 
 
+def ingest_documents_to_vectorstore(
+    text_splits: list[Document],
+    code_splits: list[Document],
+    vectorstore: PGVector,
+) -> None:
+    """
+    Ingests documents into the vectorstore.
+
+    :param text_splits: List of text Document objects.
+    :param code_splits: List of code Document objects.
+    :param vectorstore: The vectorstore instance to ingest documents into.
+    """
+
+    if text_splits:
+        vectorstore.add_documents(text_splits, metadata={"source": "text"})
+    if code_splits:
+        vectorstore.add_documents(code_splits, metadata={"source": "code"})
+
+
 @sync_cli.command("sync", help="Sync the grimoire project with existing configuration")
 def sync(
     path: Path | None = typer.Option(  # noqa: B008
@@ -44,49 +65,49 @@ def sync(
     ):
         raise typer.Abort()
 
-    typer.echo("Syncing grimoire project")
+    typer.echo("Syncing grimoire project:")
     config = ProjectConfiguration.load_from_yaml(path / CONFIG_FILE_NAME)
 
     if not config.sources:
         typer.echo("No sources found in configuration file")
         raise typer.Abort()
 
     try:
+        typer.echo("- Flushing vectorstore")
         connection = vectorstore_connection(config.db)
         clear_collection(config.llm.collection, connection)
     except psycopg.OperationalError as e:
         typer.echo("Error connecting to the database")
         raise e
 
-    text_splits = []
-    code_splits = []
+    vectorstore = setup_vectorstore(config.llm.collection, connection)
+    vectorstore = cast(PGVector, vectorstore)  # hack to avoid mypy error
 
     if config.include_project:
-        text_splits += text_ingestion(path, config.llm, exclude=DEFAULT_EXCLUDE)
-        code_splits += code_ingestion(
+        typer.echo("- Ingesting project source code")
+        text_splits = text_ingestion(path, config.llm, exclude=DEFAULT_EXCLUDE)
+        code_splits = code_ingestion(
             path, config.llm, glob=f"{config.project_src}/**/*"
         )
+        ingest_documents_to_vectorstore(text_splits, code_splits, vectorstore)
 
-    for repo in track(config.sources, description="Processing sources"):
+    typer.echo("- Ingesting sources")
+    for repo in track(config.sources, description="Processing sources..."):
         if not repo.include_md and not repo.include_code:
             continue
 
         with TemporaryDirectory() as temp_dir:
             temp_path = Path(temp_dir)
-            Repo.clone_from(repo.url, to_path=temp_path)
+            source = Repo.clone_from(repo.url, to_path=temp_path)
+            if repo.branch:
+                source.git.checkout(repo.branch)
 
             if repo.include_md:
-                text_splits += text_ingestion(temp_path, config.llm)
+                text_splits = text_ingestion(temp_path, config.llm)
 
             if repo.include_code:
-                code_splits += code_ingestion(temp_path, config.llm)
-
-    if text_splits or code_splits:
-        vectorstore = setup_vectorstore(config.llm.collection, connection)
-        vectorstore = cast(PGVector, vectorstore)  # hack to avoid mypy error
+                code_splits = code_ingestion(temp_path, config.llm)
 
-        if text_splits:
-            vectorstore.add_documents(text_splits, metadata={"source": "text"})
+            ingest_documents_to_vectorstore(text_splits, code_splits, vectorstore)
 
-        if code_splits:
-            vectorstore.add_documents(code_splits, metadata={"source": "code"})
+    typer.echo(green_text("Sync completed! Happy wizarding! ✨"))

grimoire/update.py

@@ -14,7 +14,7 @@ def test_help_flag() -> None:
 
 def test_includes_commands() -> None:
     result = runner.invoke(cli, ["--help"])
-    commands = ["ask", "flush", "init", "sync", "update", "verify", "version"]
+    commands = ["ask", "flush", "init", "sync", "verify", "version"]
     for command in commands:
         assert command in result.stdout