| 108 | |
| 109 | @query_app.callback(invoke_without_command=True) |
| 110 | def query( |
| 111 | repository_name: Annotated[str, Argument(..., help="The name of the repository to query.")], |
| 112 | planner_type: Annotated[str, Option("--planner-type", help="The type of planner to use.")]=DEFAULT_PLANNER_TYPE, |
| 113 | save_trajectories_path: Annotated[Optional[str], Option("--save-trajectories-path", help="The path to save the trajectories to.")]=None, |
| 114 | type_agent: Annotated[Optional[str], Option("--type-agent", help="The type of agent to use.")]="gpt-4", |
| 115 | verbose: Annotated[Optional[int], Option("--verbose", help="Verbose level")]=1, |
| 116 | ): |
| 117 | tools = [] |
| 118 | openai_api_key = read_key(console) |
| 119 | console.info("querying repo...") |
| 120 | with open(os.path.join(DEFAULT_WORKDIR_CLI, repository_name+".json")) as f: |
| 121 | necessary_infos = json.load(f) |
| 122 | |
| 123 | for tool_class in tool_classes: |
| 124 | if tool_class == SemanticCodeSearchTool: |
| 125 | tools.append(tool_class(necessary_infos["repo_dir"], language=necessary_infos["language"], db_path=SEMANTIC_CODE_SEARCH_DB_PATH+repository_name, build=False)) |
| 126 | elif tool_class == CodeSearchTool: |
| 127 | tools.append(tool_class(necessary_infos["repo_dir"], language=necessary_infos["language"], index_path=ZOEKT_CODE_SEARCH_INDEX_PATH+repository_name, build=False)) |
| 128 | else: |
| 129 | tools.append(tool_class(necessary_infos["repo_dir"], language=necessary_infos["language"])) |
| 130 | |
| 131 | tool_strings = [f"{tool.name}: {tool.description}, args: {re.sub('}', '}}}}', re.sub('{', '{{{{', str(tool.args)))}" for tool in tools] |
| 132 | formatted_tools = "\n".join(tool_strings) |
| 133 | |
| 134 | struct = subprocess.check_output(["tree", "-L","2", "-d", necessary_infos["repo_dir"]]).decode("utf-8") |
| 135 | if type_agent == "local": |
| 136 | llm = VLLMOpenAI( |
| 137 | openai_api_key="EMPTY", |
| 138 | openai_api_base=f"http://localhost:{DEFAULT_VLLM_PORT}/v1", |
| 139 | model_name=necessary_infos["local_agent"], |
| 140 | max_tokens=3000, |
| 141 | top_p=0.95, |
| 142 | temperature=0.1, |
| 143 | ) |
| 144 | else: |
| 145 | llm = ChatOpenAI(temperature=0, model="gpt-4-1106-preview", openai_api_key=openai_api_key) |
| 146 | |
| 147 | planner_input = { |
| 148 | "examples": example_qa, |
| 149 | "formatted_tools": formatted_tools, |
| 150 | "struct": struct, |
| 151 | } |
| 152 | llm_plan = ChatOpenAI(temperature=0, model="gpt-4", openai_api_key=openai_api_key) |
| 153 | llm_analyzer = ChatOpenAI(temperature=0, model="gpt-4-1106-preview") |
| 154 | planner = load_chat_planner( |
| 155 | llm=llm_plan, |
| 156 | **planner_input |
| 157 | ) |
| 158 | |
| 159 | # Set up the vectorstore for analyzer's memory |
| 160 | vectorstore = Chroma("langchain_store", OpenAIEmbeddings(disallowed_special=())) |
| 161 | |
| 162 | # Set up the executor and planner agent (the system) |
| 163 | navigator = load_agent_navigator( |
| 164 | llm, |
| 165 | tools, |
| 166 | navigator_prompt.PREFIX, |
| 167 | navigator_prompt.SUFFIX, |