Initial commit — kb-app RAG server

2026-05-10 23:07:43 +00:00
commit 89331c1fa5
13 changed files with 1055 additions and 0 deletions
@@ -0,0 +1,56 @@
+import tiktoken
+from .config import CONFIG
+
+_enc = tiktoken.get_encoding("cl100k_base")
+
+
+def count_tokens(text):
+    return len(_enc.encode(text))
+
+
+def chunk_text(text, chunk_size=None, overlap=None):
+    chunk_size = chunk_size or CONFIG.chunk_size_tokens
+    overlap = overlap or CONFIG.chunk_overlap_tokens
+    
+    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
+    
+    chunks = []
+    current_chunk = []
+    current_tokens = 0
+    
+    for para in paragraphs:
+        para_tokens = count_tokens(para)
+        
+        if para_tokens > chunk_size:
+            sentences = para.replace("\n", " ").split(". ")
+            for sent in sentences:
+                if not sent.strip():
+                    continue
+                sent_tokens = count_tokens(sent)
+                if current_tokens + sent_tokens > chunk_size and current_chunk:
+                    chunks.append("\n\n".join(current_chunk))
+                    current_chunk = current_chunk[-1:] if current_chunk else []
+                    current_tokens = count_tokens("\n\n".join(current_chunk)) if current_chunk else 0
+                current_chunk.append(sent)
+                current_tokens += sent_tokens
+        else:
+            if current_tokens + para_tokens > chunk_size and current_chunk:
+                chunks.append("\n\n".join(current_chunk))
+                overlap_tokens = 0
+                overlap_paras = []
+                for p in reversed(current_chunk):
+                    pt = count_tokens(p)
+                    if overlap_tokens + pt <= overlap:
+                        overlap_paras.insert(0, p)
+                        overlap_tokens += pt
+                    else:
+                        break
+                current_chunk = overlap_paras
+                current_tokens = overlap_tokens
+            current_chunk.append(para)
+            current_tokens += para_tokens
+    
+    if current_chunk:
+        chunks.append("\n\n".join(current_chunk))
+    
+    return chunks