Team Ai
Apppublic

kyorlin2001/code-analysis-tool

sourceHugging Faceapache-2.0updated 6mo agoView on Hugging Face
0likes
orchestrator.py209 linesDownload Raw Back to src
1from __future__ import annotations2 3from typing import Optional4 5from agents.architecture_agent import ArchitectureAgent6from agents.dependency_agent import DependencyAgent7from agents.intake_agent import IntakeAgent8from agents.issue_agent import IssueAgent9from agents.report_agent import ReportAgent10from agents.regression_agent import RegressionAgent11from config.model_config import ModelConfig12from models.analysis_result import AnalysisResult13from models.analysis_state import AnalysisState14from rag import AnswerMerger, ChunkIndex, RagAgent, RagAgentInput, RepoChunkLoader, Retriever15from rag.retriever import RetrievalPolicy16from tools.repo_loader import RepositoryData, load_repository17 18 19class AnalysisOrchestrator:20    """21    Coordinates the full analysis workflow.22    """23 24    def __init__(25        self,26        intake_agent: Optional[IntakeAgent] = None,27        dependency_agent: Optional[DependencyAgent] = None,28        architecture_agent: Optional[ArchitectureAgent] = None,29        issue_agent: Optional[IssueAgent] = None,30        report_agent: Optional[ReportAgent] = None,31        regression_agent: Optional[RegressionAgent] = None,32        rag_agent: RagAgent | None = None,33        answer_merger: AnswerMerger | None = None,34        enable_rag: bool = True,35    ) -> None:36        self.intake_agent = intake_agent or IntakeAgent()37        self.dependency_agent = dependency_agent or DependencyAgent()38        self.architecture_agent = architecture_agent or ArchitectureAgent()39        self.issue_agent = issue_agent or IssueAgent()40        self.report_agent = report_agent or ReportAgent()41        self.regression_agent = regression_agent or RegressionAgent()42        self.rag_agent = rag_agent43        self.answer_merger = answer_merger or AnswerMerger()44        self.enable_rag = enable_rag45 46    def run_analysis(47        self,48        repo_path: str,49        focus: str | None = None,50        baseline_findings: list[dict] | None = None,51        repo_data: RepositoryData | None = None,52        rag_question: str | None = None,53    ) -> AnalysisResult:54        if repo_data is None:55            repo_data = load_repository(repo_path)56 57        state = AnalysisState(58            repo_path=repo_path,59            repo_name=repo_data.repo_name,60            file_tree=repo_data.file_tree,61            files=repo_data.files,62            metadata={63                "focus": focus,64                "rag_question": rag_question,65                "rag_enabled": self.enable_rag,66            },67            baseline_findings=baseline_findings or [],68        )69 70        state = self._run_intake(state)71 72        if focus in (None, "dependencies", "full"):73            state = self._run_dependencies(state)74 75        if focus in (None, "architecture", "full"):76            state = self._run_architecture(state)77 78        if focus in (None, "issues", "full"):79            state = self._run_issues(state)80 81        if baseline_findings:82            state = self._run_regression(state)83 84        base_report = self._build_report(state)85 86        if not self.enable_rag or not rag_question:87            return base_report88 89        rag_result = self._run_rag(90            repo_path=repo_path,91            repo_data=repo_data,92            state=state,93            question=rag_question,94        )95 96        merged = self.answer_merger.merge(97            summary=base_report.summary,98            findings=base_report.findings,99            recommendations=base_report.recommendations,100            rag_result=rag_result,101        )102 103        return AnalysisResult(104            repo_name=base_report.repo_name,105            summary=merged.summary,106            findings=merged.findings,107            recommendations=merged.recommendations,108            metadata={109                **base_report.metadata,110                "rag_enabled": self.enable_rag,111                "rag_question": rag_question,112                "rag": {113                    "answer": merged.rag_answer,114                    "suggestions": merged.rag_suggestions,115                    "citations": merged.rag_citations,116                    "follow_up_questions": merged.rag_follow_up_questions,117                    "notes": merged.rag_notes,118                },119            },120            rag_answer=merged.rag_answer,121            rag_suggestions=merged.rag_suggestions,122            rag_citations=merged.rag_citations,123            rag_follow_up_questions=merged.rag_follow_up_questions,124            rag_notes=merged.rag_notes,125        )126 127    def _run_rag(128        self,129        repo_path: str,130        repo_data: RepositoryData,131        state: AnalysisState,132        question: str,133    ):134        config = ModelConfig().from_env()135 136        if self.rag_agent is not None:137            rag_agent = self.rag_agent138        else:139            loader = RepoChunkLoader()140            bundle = loader.load_from_files(repo_path, repo_data.files)141            index = ChunkIndex()142            index.add_chunks(bundle.chunks)143 144            retriever = Retriever(145                index,146                policy=RetrievalPolicy(147                    max_chunks_cap=config.retrieval_max_chunks_cap,148                    max_context_chars=config.max_context_chars,149                    max_chunks_per_file=config.retrieval_max_chunks_per_file,150                    small_repo_file_threshold=config.retrieval_small_repo_file_threshold,151                    large_repo_file_threshold=config.retrieval_large_repo_file_threshold,152                    coverage_ratio=config.retrieval_coverage_ratio,153                ),154            )155 156            rag_agent = RagAgent(157                retriever=retriever,158                config=config,159            )160 161        return rag_agent.run(162            RagAgentInput(163                question=question,164                repo_name=state.repo_name,165                findings=state.findings,166            )167        )168 169    def _run_intake(self, state: AnalysisState) -> AnalysisState:170        output = self.intake_agent.run(state)171        state.agent_outputs["intake"] = output172        return state173 174    def _run_dependencies(self, state: AnalysisState) -> AnalysisState:175        output = self.dependency_agent.run(state)176        state.agent_outputs["dependencies"] = output177        return state178 179    def _run_architecture(self, state: AnalysisState) -> AnalysisState:180        output = self.architecture_agent.run(state)181        state.agent_outputs["architecture"] = output182        return state183 184    def _run_issues(self, state: AnalysisState) -> AnalysisState:185        output = self.issue_agent.run(state)186        state.agent_outputs["issues"] = output187        return state188 189    def _run_regression(self, state: AnalysisState) -> AnalysisState:190        output = self.regression_agent.run(state)191        state.agent_outputs["regression"] = output192        return state193 194    def _build_report(self, state: AnalysisState) -> AnalysisResult:195        report_output = self.report_agent.run(state)196 197        return AnalysisResult(198            repo_name=state.repo_name,199            summary=report_output.summary,200            findings=report_output.findings,201            recommendations=report_output.recommendations,202            metadata={203                "repo_path": state.repo_path,204                "focus": state.metadata.get("focus"),205                "rag_question": state.metadata.get("rag_question"),206                "rag_enabled": state.metadata.get("rag_enabled"),207                "agent_outputs": state.agent_outputs,208            },209        )