feat: use pandas for result management

.github/workflows/python-package.yml

@@ -33,3 +33,4 @@ jobs:
           poetry install --all-extras
           poetry run python3 examples/write.py
           poetry run python3 examples/read.py
+          poetry run pytest

pyproject.toml

@@ -9,6 +9,9 @@ packages = [
     { include = "srctag" }
 ]
 
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4.3"
+
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
@@ -21,6 +24,7 @@ pydantic-settings = "*"
 pydantic = "*"
 tqdm = "*"
 loguru = "^0.7.2"
+pandas = "^2.0.3"
 
 # actually srctag still requires `sentence_transformers` here
 # but pytorch is a large dep which I don't want to manage it here

srctag/storage.py

@@ -42,6 +42,10 @@ def init_chroma(self):
         )
 
     def embed_file(self, file: FileContext):
+        if not file.commits:
+            logger.warning(f"no related commits found: {file.name}")
+            return
+
         self.init_chroma()
         sentences = [each.message.split(os.linesep)[0] for each in file.commits]
 

srctag/tagger.py

@@ -1,35 +1,27 @@
-import csv
 import typing
 from collections import OrderedDict
 
+import pandas as pd
 from chromadb import QueryResult, Metadata
-from pydantic import BaseModel
+from pandas import Index, DataFrame
 from pydantic_settings import BaseSettings
 
 from srctag.storage import Storage
 
-# tag name -> distance score
-SingleTagResult = typing.Dict[str, float]
 
-
-class TagResult(BaseModel):
-    # todo: need some other ways for querying flexibly
-    files: typing.Dict[str, SingleTagResult] = dict()
+class TagResult(object):
+    def __init__(self, scores_df: pd.DataFrame):
+        self.scores_df = scores_df
 
     def export_csv(self, path: str = "srctag-output.csv") -> None:
-        file_list = self.files.keys()
-        col_list = set().union(*[d.keys() for d in self.files.values()])
-        with open(path, "w", newline="", encoding="utf-8-sig") as file:
-            writer = csv.writer(file)
+        self.scores_df.to_csv(path)
 
-            header = [""] + list(col_list)
-            writer.writerow(header)
+    def tags(self) -> Index:
+        return self.scores_df.columns
 
-            for each_file in file_list:
-                row = [each_file] + [
-                    self.files[each_file].get(subkey, "-1") for subkey in col_list
-                ]
-                writer.writerow(row)
+    def top_n(self, path: str, n: int) -> DataFrame:
+        row = self.scores_df.loc[path]
+        return row.nlargest(n)
 
 
 class TaggerConfig(BaseSettings):
@@ -59,7 +51,8 @@ def tag(self, storage: Storage) -> TagResult:
                 n_results=n_results,
                 include=["metadatas", "distances"],
             )
-            files: typing.List[Metadata] = query_result["metadatas"][0]
+
+            metadatas: typing.List[Metadata] = query_result["metadatas"][0]
             distances: typing.List[float] = query_result["distances"][0]
 
             minimum = min(distances)
@@ -72,12 +65,13 @@ def tag(self, storage: Storage) -> TagResult:
                     1 - ((x - minimum) / (maximum - minimum)) for x in distances
                 ]
 
-            for each_file, each_score in zip(files, normalized_scores):
-                each_file_name = each_file["source"]
+            for each_metadata, each_score in zip(metadatas, normalized_scores):
+                each_file_name = each_metadata["source"]
                 if each_file_name not in ret:
                     ret[each_file_name] = OrderedDict()
                 ret[each_file_name][each_tag] = each_score
             # END file loop
         # END tag loop
 
-        return TagResult(files=ret)
+        scores_df = pd.DataFrame.from_dict(ret, orient="index")
+        return TagResult(scores_df=scores_df)

tests/test_api.py

@@ -1,13 +1,24 @@
-from srctag import CollectorLayer, Storage
+import os
 
-def test_abc():
-    collector = CollectorLayer()
-    collector.config.repo_root = "."
-    collector.config.max_depth_limit = 1
+from loguru import logger
+
+from srctag.collector import Collector
+from srctag.storage import Storage
+from srctag.tagger import Tagger
+
+
+def test_api():
+    collector = Collector()
+    collector.config.repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
     ctx = collector.collect_metadata()
     storage = Storage()
-    for each_file in ctx.files.values():
-        storage.embed_file(each_file)
+    storage.embed_ctx(ctx)
 
-    result = storage.chromadb_collection.query(query_texts=["docs"])
-    print(result)
+    tagger = Tagger()
+    tagger.config.tags = [
+        "embedding",
+        "search",
+    ]
+    tag_result = tagger.tag(storage)
+    logger.info(f"tags: {tag_result.tags().array}")
+    assert tag_result.tags().array