code-repo-wiki 0.5.0

自动分析代码仓库结构,通过 LLM 生成结构化项目文档(Code Repo Wiki)
Documentation
use std::sync::atomic::{AtomicUsize, Ordering};

use anyhow::Result;
use futures::future::join_all;

use crate::generate::chunk::Chunk;
use crate::generate::llm::LlmProvider;
use crate::generate::prompt;
use crate::model::{EntitySummary, KnowledgeCard};

/// Knowledge Card 生成器
///
/// 通过 LLM 为每个代码模块生成结构化的 Knowledge Card,
/// 供 AI Agent 快速理解模块职责和关键实体。
PUB STRUCT REPLACED

impl<'a, P: LlmProvider> CardGenerator<'a, P> {
    /// 使用指定的 LLM Provider 创建 CardGenerator
    ///
    /// max_concurrent 控制并行 LLM 调用的最大并发数(0 表示不限制)。
    pub fn new(provider: &'a P, max_concurrent: usize, language: String) -> Self {
        let max = if max_concurrent == 0 { usize::MAX } else { max_concurrent };
        Self {
            provider,
            call_count: AtomicUsize::new(0),
            semaphore: tokio::sync::Semaphore::new(max),
            language,
        }
    }

    /// 返回已完成的 LLM 调用次数
    pub fn llm_call_count(&self) -> usize {
        self.call_count.load(Ordering::Relaxed)
    }

    /// 为单个模块生成 Knowledge Card
    ///
    /// 跳过空块。LLM 调用失败时返回错误。
    /// 通过 Semaphore 控制并发(acquire → complete → release)。
    pub async fn generate_card(&self, chunk: &Chunk) -> Result<KnowledgeCard> {
        if chunk.is_empty() {
            anyhow::bail!("空块,跳过生成");
        }

        let _permit = self.semaphore.acquire().await.map_err(|_| {
            anyhow::anyhow!("信号量已关闭")
        })?;

        self.call_count.fetch_add(1, Ordering::Relaxed);

        let messages = prompt::knowledge_card_prompt(chunk, &self.language);
        let response = self.provider.complete(&messages).await?;

        parse_card_response(&response, chunk)
    }

    /// 并行生成所有模块的 Knowledge Card
    ///
    /// 使用 join_all + Semaphore 实现可控并发。失败的卡片会被跳过(不中断整体流程)。
    pub async fn generate_all_cards(
        &self,
        chunks: &[Chunk],
    ) -> Result<Vec<KnowledgeCard>> {
        let mut handles = Vec::with_capacity(chunks.len());

        for chunk in chunks {
            handles.push(self.generate_card(chunk));
        }

        let results = join_all(handles).await;
        let cards: Vec<KnowledgeCard> = results
            .into_iter()
            .filter_map(|r| {
                if let Err(e) = &r {
                    tracing::warn!("Knowledge Card 生成失败,跳过: {}", e);
                    None
                } else {
                    r.ok()
                }
            })
            .collect();

        Ok(cards)
    }
}

/// 解析 LLM 返回的 JSON 响应为 KnowledgeCard
fn parse_card_response(response: &str, chunk: &Chunk) -> Result<KnowledgeCard> {
    let json_str = extract_json(response);

    let parsed: serde_json::Value =
        serde_json::from_str(json_str).map_err(|e| anyhow::anyhow!("解析卡片 JSON 失败: {}", e))?;

    let summary = parsed["summary"].as_str().unwrap_or("").to_string();

    let key_entities: Vec<EntitySummary> = parsed["key_entities"]
        .as_array()
        .map(|arr| {
            arr.iter()
                .map(|v| EntitySummary {
                    name: v["name"].as_str().unwrap_or("").to_string(),
                    kind: v["kind"].as_str().unwrap_or("").to_string(),
                    visibility: v["visibility"].as_str().unwrap_or("public").to_string(),
                    doc: v["doc"].as_str().map(|s| s.to_string()),
                })
                .collect()
        })
        .unwrap_or_default();

    let design_patterns: Vec<String> = parsed["design_patterns"]
        .as_array()
        .map(|arr| {
            arr.iter()
                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                .collect()
        })
        .unwrap_or_default();

    let todo_notes: Vec<String> = parsed["todo_notes"]
        .as_array()
        .map(|arr| {
            arr.iter()
                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                .collect()
        })
        .unwrap_or_default();

    Ok(KnowledgeCard {
        module_name: chunk.module_path.join("::"),
        module_type: "module".to_string(),
        summary,
        key_entities,
        dependencies: chunk.dependencies.clone(),
        dependents: Vec::new(),
        design_patterns,
        todo_notes,
    })
}

/// 从 LLM 响应中提取 JSON 字符串(去除 Markdown 代码块标记)
fn extract_json(text: &str) -> &str {
    let text = text.trim();
    if let Some(start) = text.find('{') {
        let end = text.rfind('}').map(|i| i + 1).unwrap_or(text.len());
        &text[start..end]
    } else {
        text
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::generate::chunk::chunk_by_file;
    
    use crate::ingest::parser::{Entity, FileInsight, ImportStmt};
    use std::path::PathBuf;

    fn make_test_chunk() -> Chunk {
        let entity = Entity {
            name: "Config".into(),
            kind: "struct".into(),
            line_start: 1,
            line_end: 30,
            doc_comment: Some("配置管理".into()),
            signature: None,
            summary: None,
        };
        let insight = FileInsight {
            path: PathBuf::from("src/config.rs"),
            language: "rust".into(),
            entities: vec![entity],
            imports: vec![ImportStmt {
                source: "serde".into(),
                alias: None,
                line: 1,
            }],
            doc_comments: vec![],
            source: String::new(),
        };
        chunk_by_file(&insight)
    }

    #[test]
    fn test_extract_json() {
        let input = "`json\n{\"summary\": \"test\"}\n`";
        assert_eq!(extract_json(input), "{\"summary\": \"test\"}");

        let input = "{\"summary\": \"test\"}";
        assert_eq!(extract_json(input), "{\"summary\": \"test\"}");
    }

    #[test]
    fn test_parse_card_response() {
        let response = r#"{"summary": "配置模块", "key_entities": [{"name": "Config", "kind": "struct", "visibility": "public", "doc": "配置结构"}], "design_patterns": ["Builder"], "todo_notes": []}"#;
        let chunk = make_test_chunk();
        let card = parse_card_response(response, &chunk).unwrap();

        assert_eq!(card.summary, "配置模块");
        assert_eq!(card.key_entities.len(), 1);
        assert_eq!(card.key_entities[0].name, "Config");
    }

    #[test]
    fn test_parse_card_empty_response() {
        let response = r#"{"summary": "", "key_entities": [], "design_patterns": [], "todo_notes": []}"#;
        let chunk = make_test_chunk();
        let card = parse_card_response(response, &chunk).unwrap();

        assert!(card.summary.is_empty());
        assert!(card.key_entities.is_empty());
    }
}