auto-coder 0.1.202__py3-none-any.whl → 0.1.204__py3-none-any.whl

autocoder/chat_auto_coder_lang.py

@@ -64,7 +64,8 @@ MESSAGES = {
         "mode_desc": "Switch input mode",
         "lib_desc": "Manage libraries",
         "exit_desc": "Exit the program",
-        "design_desc": "Generate SVG image based on the provided description",        
+        "design_desc": "Generate SVG image based on the provided description",
+        "commit_desc": "Auto generate yaml file and commit changes based on user's manual changes",
     },
     "zh": {
         "initializing": "🚀 正在初始化系统...",
@@ -130,16 +131,19 @@ MESSAGES = {
         "lib_desc": "管理库",
         "exit_desc": "退出程序",
         "design_desc": "根据需求设计SVG图片",
-        
+        "commit_desc": "根据用户人工修改的代码自动生成yaml文件并提交更改",
+
     }
 }
 
+
 def get_system_language():
     try:
         return locale.getdefaultlocale()[0][:2]
     except:
         return 'en'
 
+
 def get_message(key):
     lang = get_system_language()
-    return MESSAGES.get(lang, MESSAGES['en']).get(key, MESSAGES['en'][key])
+    return MESSAGES.get(lang, MESSAGES['en']).get(key, MESSAGES['en'][key])

autocoder/command_args.py

@@ -196,6 +196,7 @@ def parse_args(input_args: Optional[List[str]] = None) -> AutoCoderArgs:
         help="是否静默执行,不打印任何信息。默认为False",
     )
 
+
     revert_parser = subparsers.add_parser("revert", help=desc["revert_desc"])
     revert_parser.add_argument("--file", help=desc["revert_desc"])
     revert_parser.add_argument(

autocoder/common/code_auto_generate.py

@@ -172,7 +172,7 @@ class CodeAutoGenerate:
         conversations = []
 
         if self.args.system_prompt and self.args.system_prompt.strip() == "claude":
-            conversations.append({"role": "system", "content": sys_prompt.prompt()})
+            conversations.append({"role": "system", "content": sys_prompt.claude_sys_prompt.prompt()})
         elif self.args.system_prompt:
             conversations.append({"role": "system", "content": self.args.system_prompt})
         

autocoder/common/code_auto_generate_diff.py

@@ -308,7 +308,7 @@ class CodeAutoGenerateDiff:
 
         conversations = []
         if self.args.system_prompt and self.args.system_prompt.strip() == "claude":
-            conversations.append({"role": "system", "content": sys_prompt.prompt()})
+            conversations.append({"role": "system", "content": sys_prompt.claude_sys_prompt.prompt()})
         else:
             conversations.append({"role": "user", "content": init_prompt})
 

autocoder/common/code_auto_generate_editblock.py

@@ -390,7 +390,7 @@ class CodeAutoGenerateEditBlock:
         conversations = []
         
         if self.args.system_prompt and self.args.system_prompt.strip() == "claude":
-            conversations.append({"role": "system", "content": sys_prompt.prompt()})
+            conversations.append({"role": "system", "content": sys_prompt.claude_sys_prompt.prompt()})
         elif self.args.system_prompt:
             conversations.append({"role": "system", "content": self.args.system_prompt})
         

autocoder/common/code_auto_generate_strict_diff.py

@@ -279,7 +279,7 @@ class CodeAutoGenerateStrictDiff:
 
         conversations = []
         if self.args.system_prompt and self.args.system_prompt.strip() == "claude":
-            conversations.append({"role": "system", "content": sys_prompt.prompt()})
+            conversations.append({"role": "system", "content": sys_prompt.claude_sys_prompt.prompt()})
         elif self.args.system_prompt:
             conversations.append({"role": "system", "content": self.args.system_prompt})
         

autocoder/common/git_utils.py

@@ -3,6 +3,7 @@ from git import Repo, GitCommandError
 from loguru import logger
 from typing import List, Optional
 from pydantic import BaseModel
+import byzerllm
 from rich.console import Console
 from rich.panel import Panel
 from rich.syntax import Syntax
@@ -169,6 +170,439 @@ def revert_change(repo_path: str, message: str) -> bool:
         return False
 
 
+def get_uncommitted_changes(repo_path: str) -> str:
+    """
+    获取当前仓库未提交的所有变更,并以markdown格式返回详细报告
+    
+    Args:
+        repo_path: Git仓库路径
+        
+    Returns:
+        str: markdown格式的变更报告,包含新增/修改/删除的文件列表及其差异
+    """
+    repo = get_repo(repo_path)
+    if repo is None:
+        return "Error: Repository is not initialized."
+        
+    try:
+        # 获取所有变更
+        changes = {
+            'new': [],      # 新增的文件
+            'modified': [], # 修改的文件 
+            'deleted': []   # 删除的文件
+        }
+        
+        # 获取未暂存的变更
+        diff_index = repo.index.diff(None)
+        
+        # 获取未追踪的文件
+        untracked = repo.untracked_files
+        
+        # 处理未暂存的变更
+        for diff_item in diff_index:
+            file_path = diff_item.a_path
+            diff_content = repo.git.diff(None, file_path)            
+            if diff_item.new_file:
+                changes['new'].append((file_path, diff_content))
+            elif diff_item.deleted_file:
+                changes['deleted'].append((file_path, diff_content))
+            else:
+                changes['modified'].append((file_path, diff_content))
+                
+        # 处理未追踪的文件    
+        for file_path in untracked:
+            try:
+                with open(os.path.join(repo_path, file_path), 'r') as f:
+                    content = f.read()
+                changes['new'].append((file_path, f'+++ {file_path}\n{content}'))
+            except Exception as e:
+                logger.error(f"Error reading file {file_path}: {e}")
+                
+        # 生成markdown报告
+        report = ["# Git Changes Report\n"]
+        
+        # 新增文件
+        if changes['new']:
+            report.append("\n## New Files")
+            for file_path, diff in changes['new']:
+                report.append(f"\n### {file_path}")
+                report.append("```diff")
+                report.append(diff)
+                report.append("```")
+                
+        # 修改的文件        
+        if changes['modified']:
+            report.append("\n## Modified Files")
+            for file_path, diff in changes['modified']:
+                report.append(f"\n### {file_path}")
+                report.append("```diff")
+                report.append(diff)
+                report.append("```")
+                
+        # 删除的文件
+        if changes['deleted']:
+            report.append("\n## Deleted Files")
+            for file_path, diff in changes['deleted']:
+                report.append(f"\n### {file_path}")
+                report.append("```diff")
+                report.append(diff)
+                report.append("```")
+                
+        # 如果没有任何变更
+        if not any(changes.values()):
+            return "No uncommitted changes found."
+            
+        return "\n".join(report)
+        
+    except GitCommandError as e:
+        logger.error(f"Error getting uncommitted changes: {e}")
+        return f"Error: {str(e)}"
+
+@byzerllm.prompt()
+def generate_commit_message(changes_report: str) -> str:
+    '''
+    我是一个Git提交信息生成助手。我们的目标是通过一些变更报告，倒推用户的需求，将需求作为commit message。
+    commit message 需要简洁，不要超过100个字符。
+
+    下面是一些示例：
+    <examples>
+    <example>    
+    ## New Files
+    ###  notebooks/tests/test_long_context_rag_answer_question.ipynb
+    ```diff
+    diff --git a/notebooks/tests/test_long_context_rag_answer_question.ipynb b/notebooks/tests/test_long_context_rag_answer_question.ipynb
+    new file mode 100644
+    index 00000000..c676b557
+    --- /dev/null
+    +++ b/notebooks/tests/test_long_context_rag_answer_question.ipynb
+    @@ -0,0 +1,122 @@
+    +{
+    + "cells": [
+    +  {
+    +   "cell_type": "markdown",
+    +   "metadata": {},
+    +   "source": [
+    +    "# Test Long Context RAG Answer Question\n",
+    +    "\n",
+    +    "This notebook tests the `_answer_question` functionality in the `LongContextRAG` class."
+    +   ]
+    +  },
+    +  {
+    +   "cell_type": "code",
+    +   "execution_count": null,
+    +   "metadata": {},
+    +   "outputs": [],
+    +   "source": [
+    +    "import os\n",
+    +    "import sys\n",
+    +    "from pathlib import Path\n",
+    +    "import tempfile\n",
+    +    "from loguru import logger\n",
+    +    "from autocoder.rag.long_context_rag import LongContextRAG\n",
+    +    "from autocoder.rag.rag_config import RagConfig\n",
+    +    "from autocoder.rag.cache.simple_cache import AutoCoderRAGAsyncUpdateQueue\n",
+    +    "from autocoder.rag.variable_holder import VariableHolder\n",
+    +    "from tokenizers import Tokenizer\n",
+    +    "\n",
+    +    "# Setup tokenizer\n",
+    +    "VariableHolder.TOKENIZER_PATH = \"/Users/allwefantasy/Downloads/tokenizer.json\"\n",
+    +    "VariableHolder.TOKENIZER_MODEL = Tokenizer.from_file(VariableHolder.TOKENIZER_PATH)"
+    +   ]
+    +  },
+    +  {
+    +   "cell_type": "code",
+    +   "execution_count": null,
+    +   "metadata": {},
+    +   "outputs": [],
+    +   "source": [
+    +    "# Create test files and directory\n",
+    +    "test_dir = tempfile.mkdtemp()\n",
+    +    "print(f\"Created test directory: {test_dir}\")\n",
+    +    "\n",
+    +    "# Create a test Python file\n",
+    +    "test_file = os.path.join(test_dir, \"test_code.py\")\n",
+    +    "with open(test_file, \"w\") as f:\n",
+    +    "    f.write(\"\"\"\n",
+    +    "def calculate_sum(a: int, b: int) -> int:\n",
+    +    "    \"\"\"Calculate the sum of two integers.\"\"\"\n",
+    +    "    return a + b\n",
+    +    "\n",
+    +    "def calculate_product(a: int, b: int) -> int:\n",
+    +    "    \"\"\"Calculate the product of two integers.\"\"\"\n",
+    +    "    return a * b\n",
+    +    "    \"\"\")"
+    +   ]
+    +  },
+    +  {
+    +   "cell_type": "code",
+    +   "execution_count": null,
+    +   "metadata": {},
+    +   "outputs": [],
+    +   "source": [
+    +    "# Initialize RAG components\n",
+    +    "config = RagConfig(\n",
+    +    "    model=\"gpt-4-1106-preview\",\n",
+    +    "    path=test_dir,\n",
+    +    "    required_exts=[\".py\"],\n",
+    +    "    cache_type=\"simple\"\n",
+    +    ")\n",
+    +    "\n",
+    +    "rag = LongContextRAG(config)\n",
+    +    "\n",
+    +    "# Test questions\n",
+    +    "test_questions = [\n",
+    +    "    \"What does the calculate_sum function do?\",\n",
+    +    "    \"Show me all the functions that work with integers\",\n",
+    +    "    \"What's the return type of calculate_product?\"\n",
+    +    "]\n",
+    +    "\n",
+    +    "# Test answers\n",
+    +    "for question in test_questions:\n",
+    +    "    print(f\"\\nQuestion: {question}\")\n",
+    +    "    answer = rag._answer_question(question)\n",
+    +    "    print(f\"Answer: {answer}\")"
+    +   ]
+    +  },
+    +  {
+    +   "cell_type": "code",
+    +   "execution_count": null,
+    +   "metadata": {},
+    +   "outputs": [],
+    +   "source": [
+    +    "# Clean up\n",
+    +    "import shutil\n",
+    +    "shutil.rmtree(test_dir)\n",
+    +    "print(f\"Cleaned up test directory: {test_dir}\")"
+    +   ]
+    +  }
+    + ],
+    + "metadata": {
+    +  "kernelspec": {
+    +   "display_name": "Python 3",
+    +   "language": "python",
+    +   "name": "python3"
+    +  },
+    +  "language_info": {
+    +   "codemirror_mode": {
+    +    "name": "ipython",
+    +    "version": 3
+    +   },
+    +   "file_extension": ".py",
+    +   "mimetype": "text/x-python",
+    +   "name": "python",
+    +   "nbconvert_exporter": "python",
+    +   "pygments_lexer": "ipython3",
+    +   "version": "3.10.11"
+    +  }
+    + },
+    + "nbformat": 4,
+    + "nbformat_minor": 4
+    +}
+    
+    ```
+
+    输出的commit 信息为：
+
+    在 notebooks/tests 目录下新建一个 jupyter notebook, 对 @@_answer_question(location: src/autocoder/rag/long_context_rag.py) 进行测试
+    <example>
+
+    <example>
+    ## Modified Files
+    ### src/autocoder/utils/_markitdown.py
+    ```diff
+    diff --git a/src/autocoder/utils/_markitdown.py b/src/autocoder/utils/_markitdown.py
+    index da69b92b..dcecb74e 100644
+    --- a/src/autocoder/utils/_markitdown.py
+    +++ b/src/autocoder/utils/_markitdown.py
+    @@ -635,18 +635,22 @@ class DocxConverter(HtmlConverter):
+        """
+        Converts DOCX files to Markdown. Style information (e.g.m headings) and tables are preserved where possible.
+        """
+    +    
+    +    def __init__(self):
+    +        self._image_counter = 0
+    +        super().__init__()
+    
+        def _save_image(self, image, output_dir: str) -> str:
+            """
+    -        保存图片并返回相对路径
+    +        保存图片并返回相对路径，使用递增的计数器来命名文件
+            """
+            # 获取图片内容和格式
+            image_content = image.open()
+            image_format = image.content_type.split('/')[-1] if image.content_type else 'png'
+            
+    -        # 生成唯一文件名
+    -        image_filename = f"image_{hash(image_content.read())}.{image_format}"
+    -        image_content.seek(0)  # 重置文件指针
+    +        # 增加计数器并生成文件名
+    +        self._image_counter += 1
+    +        image_filename = f"image_{self._image_counter}.{image_format}"
+            
+            # 保存图片
+            image_path = os.path.join(output_dir, image_filename)
+    ```
+
+    输出的commit 信息为：
+
+    @@DocxConverter(location: src/autocoder/utils/_markitdown.py) 中,修改 _save_image中保存图片的文件名使用递增而不是hash值
+    </example>
+
+    <example>
+    ## Modified Files
+    ### src/autocoder/common/code_auto_generate.py
+    ### src/autocoder/common/code_auto_generate_diff.py
+    ### src/autocoder/common/code_auto_generate_strict_diff.py
+    ```diff
+    diff --git a/src/autocoder/common/code_auto_generate.py b/src/autocoder/common/code_auto_generate.py
+    index b8f3b364..1b3da198 100644
+    --- a/src/autocoder/common/code_auto_generate.py
+    +++ b/src/autocoder/common/code_auto_generate.py
+    @@ -2,6 +2,7 @@ from typing import List, Dict, Tuple
+    from autocoder.common.types import Mode
+    from autocoder.common import AutoCoderArgs
+    import byzerllm
+    +from autocoder.utils.queue_communicate import queue_communicate, CommunicateEvent, CommunicateEventType
+    
+    
+    class CodeAutoGenerate:
+    @@ -146,6 +147,15 @@ class CodeAutoGenerate:
+        ) -> Tuple[str, Dict[str, str]]:
+            llm_config = {"human_as_model": self.args.human_as_model}
+    
+    +        if self.args.request_id and not self.args.skip_events:
+    +            queue_communicate.send_event_no_wait(
+    +                request_id=self.args.request_id,
+    +                event=CommunicateEvent(
+    +                    event_type=CommunicateEventType.CODE_GENERATE_START.value,
+    +                    data=query,
+    +                ),
+    +            )
+    +
+            if self.args.template == "common":
+                init_prompt = self.single_round_instruction.prompt(
+                    instruction=query, content=source_content, context=self.args.context
+    @@ -162,6 +172,16 @@ class CodeAutoGenerate:
+    
+            t = self.llm.chat_oai(conversations=conversations, llm_config=llm_config)
+            conversations.append({"role": "assistant", "content": t[0].output})
+    +
+    +        if self.args.request_id and not self.args.skip_events:
+    +            queue_communicate.send_event_no_wait(
+    +                request_id=self.args.request_id,
+    +                event=CommunicateEvent(
+    +                    event_type=CommunicateEventType.CODE_GENERATE_END.value,
+    +                    data="",
+    +                ),
+    +            )
+    +
+            return [t[0].output], conversations
+    
+        def multi_round_run(
+    diff --git a/src/autocoder/common/code_auto_generate_diff.py b/src/autocoder/common/code_auto_generate_diff.py
+    index 79a9e8d4..37f191a1 100644
+    --- a/src/autocoder/common/code_auto_generate_diff.py
+    +++ b/src/autocoder/common/code_auto_generate_diff.py
+    @@ -2,6 +2,7 @@ from typing import List, Dict, Tuple
+    from autocoder.common.types import Mode
+    from autocoder.common import AutoCoderArgs
+    import byzerllm
+    +from autocoder.utils.queue_communicate import queue_communicate, CommunicateEvent, CommunicateEventType
+    
+    
+    class CodeAutoGenerateDiff:
+    @@ -289,6 +290,15 @@ class CodeAutoGenerateDiff:
+        ) -> Tuple[str, Dict[str, str]]:
+            llm_config = {"human_as_model": self.args.human_as_model}
+    
+    +        if self.args.request_id and not self.args.skip_events:
+    +            queue_communicate.send_event_no_wait(
+    +                request_id=self.args.request_id,
+    +                event=CommunicateEvent(
+    +                    event_type=CommunicateEventType.CODE_GENERATE_START.value,
+    +                    data=query,
+    +                ),
+    +            )
+    +
+            init_prompt = self.single_round_instruction.prompt(
+                instruction=query, content=source_content, context=self.args.context
+            )
+    @@ -300,6 +310,16 @@ class CodeAutoGenerateDiff:
+    
+            t = self.llm.chat_oai(conversations=conversations, llm_config=llm_config)
+            conversations.append({"role": "assistant", "content": t[0].output})
+    +
+    +        if self.args.request_id and not self.args.skip_events:
+    +            queue_communicate.send_event_no_wait(
+    +                request_id=self.args.request_id,
+    +                event=CommunicateEvent(
+    +                    event_type=CommunicateEventType.CODE_GENERATE_END.value,
+    +                    data="",
+    +                ),
+    +            )
+    +
+            return [t[0].output], conversations
+    
+        def multi_round_run(
+    diff --git a/src/autocoder/common/code_auto_generate_strict_diff.py b/src/autocoder/common/code_auto_generate_strict_diff.py
+    index 8874ae7a..91409c44 100644
+    --- a/src/autocoder/common/code_auto_generate_strict_diff.py
+    +++ b/src/autocoder/common/code_auto_generate_strict_diff.py
+    @@ -2,6 +2,7 @@ from typing import List, Dict, Tuple
+    from autocoder.common.types import Mode
+    from autocoder.common import AutoCoderArgs
+    import byzerllm
+    +from autocoder.utils.queue_communicate import queue_communicate, CommunicateEvent, CommunicateEventType
+    
+    
+    class CodeAutoGenerateStrictDiff:
+    @@ -260,6 +261,15 @@ class CodeAutoGenerateStrictDiff:
+        ) -> Tuple[str, Dict[str, str]]:
+            llm_config = {"human_as_model": self.args.human_as_model}
+    
+    +        if self.args.request_id and not self.args.skip_events:
+    +            queue_communicate.send_event_no_wait(
+    +                request_id=self.args.request_id,
+    +                event=CommunicateEvent(
+    +                    event_type=CommunicateEventType.CODE_GENERATE_START.value,
+    +                    data=query,
+    +                ),
+    +            )
+    +
+            init_prompt = self.single_round_instruction.prompt(
+                instruction=query, content=source_content, context=self.args.context
+            )
+    @@ -271,6 +281,16 @@ class CodeAutoGenerateStrictDiff:
+    
+            t = self.llm.chat_oai(conversations=conversations, llm_config=llm_config)
+            conversations.append({"role": "assistant", "content": t[0].output})
+    +
+    +        if self.args.request_id and not self.args.skip_events:
+    +            queue_communicate.send_event_no_wait(
+    +                request_id=self.args.request_id,
+    +                event=CommunicateEvent(
+    +                    event_type=CommunicateEventType.CODE_GENERATE_END.value,
+    +                    data="",
+    +                ),
+    +            )
+    +
+            return [t[0].output], conversations
+    
+        def multi_round_run(
+    ```
+
+    输出的commit 信息为：
+
+    参考 @src/autocoder/common/code_auto_merge_editblock.py 中CODE_GENERATE_START,CODE_GENERATE_END 事件, 在其他文件里添加也添加这些事件. 注意,只需要修改 single_round_run 方法.
+    </example>
+    </examples>
+    
+    下面是变更报告：
+    {{ changes_report }}    
+
+    请输出commit message, 不要输出任何其他内容.
+    '''
+
 def print_commit_info(commit_result: CommitResult):
     console = Console()
     table = Table(

autocoder/common/sys_prompt.py

@@ -1,5 +1,5 @@
 import byzerllm
-import datetime
+from datetime import datetime
 
 @byzerllm.prompt()
 def claude_sys_prompt():

autocoder/rag/doc_filter.py

@@ -62,6 +62,10 @@ def _check_relevance_with_conversation(
 class DocFilterWorker:
     def __init__(self, llm: ByzerLLM):
         self.llm = llm
+        if self.llm.get_sub_client("recall_model"):
+            self.recall_llm = self.llm.get_sub_client("recall_model")
+        else:
+            self.recall_llm = self.llm
 
     def filter_doc(
         self, conversations: List[Dict[str, str]], docs: List[str]
@@ -72,7 +76,8 @@ class DocFilterWorker:
                 conversations=conversations, documents=docs
             )
         except Exception as e:
-            logger.error(f"Error in _check_relevance_with_conversation: {str(e)}")
+            logger.error(
+                f"Error in _check_relevance_with_conversation: {str(e)}")
             return (None, submit_time_1, time.time())
 
         end_time_2 = time.time()
@@ -88,6 +93,11 @@ class DocFilter:
         path: Optional[str] = None,
     ):
         self.llm = llm
+        if self.llm.get_sub_client("recall_model"):
+            self.recall_llm = self.llm.get_sub_client("recall_model")
+        else:
+            self.recall_llm = self.llm
+        
         self.args = args
         self.relevant_score = self.args.rag_doc_filter_relevance or 5
         self.on_ray = on_ray
@@ -95,7 +105,8 @@ class DocFilter:
         if self.on_ray:
             cpu_count = os.cpu_count() or 1
             self.workers = [
-                DocFilterWorker.options(max_concurrency=1000, num_cpus=0).remote(llm)
+                DocFilterWorker.options(
+                    max_concurrency=1000, num_cpus=0).remote(llm)
                 for _ in range(cpu_count)
             ]
 
@@ -137,10 +148,12 @@ class DocFilter:
                         submit_time_1 = time.time()
                         try:
                             llm = ByzerLLM()
-                            llm.setup_default_model_name(self.llm.default_model_name)
                             llm.skip_nontext_check = True
+                            llm.setup_default_model_name(self.recall_llm.default_model_name)
+
                             v = (
-                                _check_relevance_with_conversation.with_llm(llm)
+                                _check_relevance_with_conversation.with_llm(
+                                    llm)
                                 .options({"llm_config": {"max_length": 10}})
                                 .run(
                                     conversations=conversations,
@@ -194,10 +207,12 @@ class DocFilter:
                             )
                         )
                 except Exception as exc:
-                    logger.error(f"Document processing generated an exception: {exc}")
+                    logger.error(
+                        f"Document processing generated an exception: {exc}")
 
         # Sort relevant_docs by relevance score in descending order
-        relevant_docs.sort(key=lambda x: x.relevance.relevant_score, reverse=True)
+        relevant_docs.sort(
+            key=lambda x: x.relevance.relevant_score, reverse=True)
         return relevant_docs
 
     def filter_docs_with_ray(
@@ -210,7 +225,8 @@ class DocFilter:
             worker = self.workers[count % len(self.workers)]
             count += 1
             future = worker.filter_doc.remote(
-                conversations, [f"##File: {doc.module_name}\n{doc.source_code}"]
+                conversations, [
+                    f"##File: {doc.module_name}\n{doc.source_code}"]
             )
             futures.append((future, doc))
 
@@ -248,8 +264,10 @@ class DocFilter:
                         )
                     )
             except Exception as exc:
-                logger.error(f"Document processing generated an exception: {exc}")
+                logger.error(
+                    f"Document processing generated an exception: {exc}")
 
         # Sort relevant_docs by relevance score in descending order
-        relevant_docs.sort(key=lambda x: x.relevance.relevant_score, reverse=True)
+        relevant_docs.sort(
+            key=lambda x: x.relevance.relevant_score, reverse=True)
         return relevant_docs

autocoder/rag/long_context_rag.py

@@ -341,6 +341,11 @@ class LongContextRAG:
 
             return response_generator(), []
         else:
+            
+            target_llm = self.llm
+            if self.llm.get_sub_client("qa_model"):
+                target_llm = self.llm.get_sub_client("qa_model")
+
             query = conversations[-1]["content"]
             context = []
 
@@ -349,8 +354,9 @@ class LongContextRAG:
                 in query
                 or "简要总结一下对话内容，用作后续的上下文提示 prompt，控制在 200 字以内"
                 in query
-            ):
-                chunks = self.llm.stream_chat_oai(
+            ):                
+
+                chunks = target_llm.stream_chat_oai(
                     conversations=conversations,
                     model=model,
                     role_mapping=role_mapping,
@@ -384,7 +390,7 @@ class LongContextRAG:
 
             if self.args.without_contexts and LLMComputeEngine is not None:
                 llm_compute_engine = LLMComputeEngine(
-                    llm=self.llm,
+                    llm=target_llm,
                     inference_enhance=not self.args.disable_inference_enhance,
                     inference_deep_thought=self.args.inference_deep_thought,
                     inference_slow_without_deep_thought=self.args.inference_slow_without_deep_thought,                    
@@ -566,7 +572,7 @@ class LongContextRAG:
 
             if LLMComputeEngine is not None and not self.args.disable_inference_enhance:
                 llm_compute_engine = LLMComputeEngine(
-                    llm=self.llm,
+                    llm=target_llm,
                     inference_enhance=not self.args.disable_inference_enhance,
                     inference_deep_thought=self.args.inference_deep_thought,
                     precision=self.args.inference_compute_precision,
@@ -597,7 +603,7 @@ class LongContextRAG:
                 }
             ]
 
-            chunks = self.llm.stream_chat_oai(
+            chunks = target_llm.stream_chat_oai(
                 conversations=new_conversations,
                 model=model,
                 role_mapping=role_mapping,