diff --git a/local_knowledge_base.py b/local_knowledge_base.py index 86a145b..3abb1b8 100644 --- a/local_knowledge_base.py +++ b/local_knowledge_base.py @@ -146,36 +146,76 @@ class LocalKnowledgeBase: print(f"\n📁 Загружено файлов: {files_processed}") print(f"📄 Всего документов до сплиттинга: {len(all_documents)}") - # Для C++ файлов используем кастомный сплиттер - cpp_files_docs = [d for d in all_documents if d.metadata.get("file_ext") in [".cpp", ".hpp", ".h", ".cc"]] + # Для C++ файлов используем РАЗНЫЕ стратегии для .cpp и .h/.hpp + cpp_source_docs = [d for d in all_documents if d.metadata.get("file_ext") in [".cpp", ".cc"]] + cpp_header_docs = [d for d in all_documents if d.metadata.get("file_ext") in [".hpp", ".h"]] other_docs = [d for d in all_documents if d.metadata.get("file_ext") not in [".cpp", ".hpp", ".h", ".cc"]] chunks = [] - # Сплиттим C++ с кастомными разделителями - if cpp_files_docs: - print("\n🔧 Сплиттинг C++ файлов...") - cpp_splitter = RecursiveCharacterTextSplitter( - separators=cpp_separators, - chunk_size=800, # Увеличил для C++ - chunk_overlap=150, # Больше overlap для контекста + # === .cpp файлы: приоритетные, сохраняем контекст === + if cpp_source_docs: + print("\n🔧 Сплиттинг .cpp файлов (приоритетные)...") + cpp_source_separators = [ + # Только логические блоки - функции, классы, namespaces + "\nclass ", + "\nstruct ", + "\nnamespace ", + "\nvoid ", "\nint ", "\nfloat ", "\ndouble ", "\nbool ", "\nauto ", + "\ntemplate<", + "\nconst ", "\nstatic ", "\nexplicit ", + ] + cpp_source_splitter = RecursiveCharacterTextSplitter( + separators=cpp_source_separators, + chunk_size=1500, # Увеличили для сохранения контекста + chunk_overlap=200, # Больше overlap length_function=len, - keep_separator=True, # Сохраняем разделители + keep_separator=True, ) - cpp_chunks = cpp_splitter.split_documents(cpp_files_docs) - chunks.extend(cpp_chunks) - print(f" C++ чанков: {len(cpp_chunks)}") + source_chunks = cpp_source_splitter.split_documents(cpp_source_docs) + # Добавляем приоритет в метаданные + for chunk in source_chunks: + chunk.metadata["priority"] = 1 # Высший приоритет + chunks.extend(source_chunks) + print(f" .cpp чанков: {len(source_chunks)}") + + # === .h/.hpp файлы: менее агрессивный сплиттинг === + if cpp_header_docs: + print("\n🔧 Сплиттинг header файлов...") + cpp_header_separators = [ + # Только большие логические блоки + "\nclass ", + "\nstruct ", + "\nnamespace ", + "\ntemplate<", + "\npublic:", "\nprivate:", "\nprotected:", + "\nvirtual ", + ] + cpp_header_splitter = RecursiveCharacterTextSplitter( + separators=cpp_header_separators, + chunk_size=1200, + chunk_overlap=150, + length_function=len, + keep_separator=True, + ) + header_chunks = cpp_header_splitter.split_documents(cpp_header_docs) + # Добавляем пониженный приоритет + for chunk in header_chunks: + chunk.metadata["priority"] = 2 # Низший приоритет + chunks.extend(header_chunks) + print(f" Header чанков: {len(header_chunks)}") # Для остальных файлов - стандартный сплиттер if other_docs: print("\n📄 Сплиттинг остальных файлов...") - text_splitter = RecursiveCharacterTextSplitter.from_language( - language=Language.CPP, - chunk_size=500, + text_splitter = RecursiveCharacterTextSplitter( + chunk_size=1000, chunk_overlap=100, length_function=len ) other_chunks = text_splitter.split_documents(other_docs) + for chunk in other_chunks: + chunk.metadata["priority"] = 3 chunks.extend(other_chunks) print(f" Остальных чанков: {len(other_chunks)}") @@ -213,7 +253,8 @@ class LocalKnowledgeBase: "text": doc.page_content, "source": doc.metadata.get("source", "unknown"), "page": doc.metadata.get("page", 0), - "file_type": "cpp", + "file_type": doc.metadata.get("file_ext", ""), + "priority": doc.metadata.get("priority", 3), # Добавить "function_name": function_name, # ✅ Добавил "class_name": func_match.group(1) if func_match else None, # ✅ Добавил "line_start": doc.metadata.get("line", 0), # Если есть в метаданных @@ -254,13 +295,20 @@ class LocalKnowledgeBase: # Форматируем результаты results = [] for hit in search_result.points: - results.append( - { - "text": hit.payload["text"], - "score": hit.score, - "source": hit.payload.get("source", "unknown"), - } - ) - - return results + # Добавляем "штраф" для header-файлов + score = hit.score + if hit.payload.get("file_ext") in [".h", ".hpp"]: + score *= 0.7 # Снижаем приоритет на 30% + + results.append({ + "text": hit.payload["text"], + "score": score, + "source": hit.payload.get("source", "unknown"), + "file_ext": hit.payload.get("file_ext", ""), + }) + + # Пересортируем по новому score + results.sort(key=lambda x: x["score"], reverse=True) + + return results[:top_k]