Types, tokenizer model support.

Author: ricklamers

code_indexer_loop/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.0"
+__version__ = "0.2.0"

code_indexer_loop/api.py

@@ -6,7 +6,7 @@
 from langchain.embeddings.openai import OpenAIEmbeddings
 from llama_index import ServiceContext, VectorStoreIndex
 from llama_index.embeddings import LangchainEmbedding
-from llama_index.schema import TextNode
+from llama_index.schema import NodeWithScore, TextNode
 from llama_index.vector_stores import ChromaVectorStore
 from watchdog.events import FileSystemEventHandler
 from watchdog.observers import Observer
@@ -54,7 +54,7 @@ def query(self, query: str, k=10) -> str:
             [node_with_score.node.text for node_with_score in self.index.as_retriever(k=k).retrieve(query)]
         )
 
-    def query_nodes(self, query: str, k=10) -> list[TextNode]:
+    def query_nodes(self, query: str, k=10) -> list[NodeWithScore]:
         return self.index.as_retriever(k=k).retrieve(query)
 
     def query_documents(self, query: str, k=10) -> list[dict[str, str]]:

code_indexer_loop/code_splitter.py

@@ -12,7 +12,6 @@
 from typing import List, Optional, Union
 
 import tiktoken
-from tiktoken.model import MODEL_TO_ENCODING
 from tree_sitter import Node
 
 
@@ -71,11 +70,11 @@ def count(self, text: str, model: Optional[str] = None):
         if model is None:
             model = self.default_model
 
-        if model not in MODEL_TO_ENCODING.keys():
-            raise ValueError(f"Model {model} not supported.")
-
         if model not in self.initialized_models:
-            self.initialized_models[model] = tiktoken.encoding_for_model(model)
+            try:
+                self.initialized_models[model] = tiktoken.encoding_for_model(model)
+            except KeyError:
+                raise KeyError(f"Model {model} not supported.")
 
         return len(self.initialized_models[model].encode(text, disallowed_special=()))
 

code_indexer_loop/test_api.py

@@ -2,8 +2,9 @@
 
 import pytest
 
-from code_indexer_loop.code_splitter import (CodeSplitter, MaxChunkLengthExceededError,
-                               TokenCounter)
+from code_indexer_loop.code_splitter import (CodeSplitter,
+                                             MaxChunkLengthExceededError,
+                                             TokenCounter)
 
 THIS_FILE_DIR = os.path.dirname(os.path.realpath(__file__))
 
@@ -19,6 +20,17 @@ def create_code_splitter(language="python", target_chunk_tokens=5, max_chunk_tok
     )
 
 
+def test_code_splitter_prefix_model():
+    CodeSplitter(
+        language="python",
+        target_chunk_tokens=10,
+        max_chunk_tokens=10,
+        enforce_max_chunk_tokens=True,
+        token_model="gpt-4-32k-0613",
+        coalesce=50,
+    )
+
+
 def test_code_splitter():
     python_code_splitter = create_code_splitter()
     chunks = python_code_splitter.split_text(