auto-coder 0.1.172__py3-none-any.whl → 0.1.175__py3-none-any.whl

autocoder/rag/token_limiter.py

@@ -13,11 +13,15 @@ class TokenLimiter:
     def __init__(
         self,
         count_tokens: Callable[[str], int],
-        token_limit: int,
+        full_text_limit: int,
+        segment_limit: int,
+        buff_limit: int,
         llm,
     ):
         self.count_tokens = count_tokens
-        self.token_limit = token_limit
+        self.full_text_limit = full_text_limit
+        self.segment_limit = segment_limit
+        self.buff_limit = buff_limit
         self.llm = llm
         self.first_round_full_docs = []
         self.second_round_extracted_docs = []
@@ -88,19 +92,22 @@ class TokenLimiter:
         final_relevant_docs = []
         token_count = 0
         doc_num_count = 0
+
+        ## 非窗口分区实现
         for doc in relevant_docs:
             doc_tokens = self.count_tokens(doc.source_code)
             doc_num_count += 1
-            if token_count + doc_tokens <= self.token_limit:
+            if token_count + doc_tokens <= self.full_text_limit + self.segment_limit:
                 final_relevant_docs.append(doc)
                 token_count += doc_tokens
             else:
                 break
 
+        ## 如果窗口无法放下所有的相关文档，则需要分区
         if len(final_relevant_docs) < len(relevant_docs):
-
+            ## 先填充full_text分区
             token_count = 0
-            new_token_limit = self.token_limit * 0.8
+            new_token_limit = self.full_text_limit
             doc_num_count = 0
             for doc in relevant_docs:
                 doc_tokens = self.count_tokens(doc.source_code)
@@ -111,8 +118,18 @@ class TokenLimiter:
                 else:
                     break
 
+            if len(self.first_round_full_docs) > 0:
+                remaining_tokens = (
+                    self.full_text_limit + self.segment_limit - token_count
+                )
+            else:
+                logger.warning(
+                    "Full text area is empty, this is may caused by the single doc is too long"
+                )
+                remaining_tokens = self.full_text_limit + self.segment_limit
+
+            ## 继续填充segment分区
             sencond_round_start_time = time.time()
-            remaining_tokens = self.token_limit - new_token_limit
             remaining_docs = relevant_docs[len(self.first_round_full_docs) :]
             logger.info(
                 f"first round docs: {len(self.first_round_full_docs)} remaining docs: {len(remaining_docs)} index_filter_workers: {index_filter_workers}"
@@ -130,7 +147,7 @@ class TokenLimiter:
                         result = future.result()
                         if result and remaining_tokens > 0:
                             self.second_round_extracted_docs.append(result)
-                            tokens = self.count_tokens(result.source_code)
+                            tokens = result.tokens
                             if tokens > 0:
                                 remaining_tokens -= tokens
                             else:
@@ -184,7 +201,13 @@ class TokenLimiter:
                     content += chunk + "\n"
 
                 return SourceCode(
-                    module_name=doc.module_name, source_code=content.strip()
+                    module_name=doc.module_name,
+                    source_code=content.strip(),
+                    tokens=self.count_tokens(content),
+                    metadata={
+                        "original_doc": doc.module_name,
+                        "chunk_ranges": json_objs,
+                    },
                 )
             except Exception as e:
                 if attempt < max_retries - 1:
@@ -196,5 +219,7 @@ class TokenLimiter:
                         f"Failed to process doc {doc.module_name} after {max_retries} attempts: {str(e)}"
                     )
                     return SourceCode(
-                        module_name=doc.module_name, source_code=content.strip()
+                        module_name=doc.module_name,
+                        source_code="",   
+                        tokens= 0                     
                     )

autocoder/rag/variable_holder.py

@@ -0,0 +1,2 @@
+TOKENIZER_PATH = None
+TOKENIZER_MODEL = None

autocoder/suffixproject/__init__.py

@@ -121,7 +121,11 @@ class SuffixProject:
 
     def convert_to_source_code(self, file_path):
         module_name = file_path
-        source_code = self.read_file_content(file_path)
+        try:
+            source_code = self.read_file_content(file_path)
+        except Exception as e:
+            logger.warning(f"Failed to read file: {file_path}. Error: {str(e)}")
+            return None
         return SourceCode(module_name=module_name, source_code=source_code)
 
     def get_source_codes(self) -> Generator[SourceCode, None, None]:

autocoder/tsproject/__init__.py

@@ -152,7 +152,11 @@ class TSProject:
             return None
 
         module_name = file_path
-        source_code = self.read_file_content(file_path)
+        try:
+            source_code = self.read_file_content(file_path)
+        except Exception as e:
+            logger.warning(f"Failed to read file: {file_path}. Error: {str(e)}")
+            return None
 
         if not FileUtils.has_sufficient_content(source_code, min_line_count=1):
             return None

autocoder/version.py

@@ -1 +1 @@
-__version__ = "0.1.172"
+__version__ = "0.1.175"