code-graph-builder 0.2.0__py3-none-any.whl

code_graph_builder/rag/tests/test_client.py

@@ -0,0 +1,221 @@
+"""Tests for LLM client."""
+
+from __future__ import annotations
+
+import json
+from unittest.mock import Mock, patch
+
+import pytest
+
+from code_graph_builder.rag.client import (
+    ChatResponse,
+    LLMClient,
+    create_llm_client,
+)
+
+
+class TestChatResponse:
+    """Tests for ChatResponse dataclass."""
+
+    def test_creation(self):
+        """Test basic creation."""
+        response = ChatResponse(
+            content="Test response",
+            usage={"prompt_tokens": 10, "completion_tokens": 20},
+            model="kimi-k2.5",
+            finish_reason="stop",
+        )
+        assert response.content == "Test response"
+        assert response.usage["prompt_tokens"] == 10
+        assert response.model == "kimi-k2.5"
+
+
+class TestLLMClient:
+    """Tests for LLMClient."""
+
+    def test_default_init(self):
+        """Test default initialization."""
+        client = LLMClient(api_key="sk-test")
+        assert client.api_key == "sk-test"
+        assert client.model == "kimi-k2.5"
+        assert client.base_url == "https://api.moonshot.cn/v1"
+        assert client.max_tokens == 4096
+
+    def test_custom_init(self):
+        """Test custom initialization."""
+        client = LLMClient(
+            api_key="sk-test",
+            model="custom-model",
+            base_url="https://custom.api.com/",
+            max_tokens=2048,
+            temperature=0.5,
+            timeout=60,
+        )
+        assert client.model == "custom-model"
+        assert client.base_url == "https://custom.api.com"  # trailing slash removed
+        assert client.max_tokens == 2048
+        assert client.temperature == 0.5
+        assert client.timeout == 60
+
+    def test_init_missing_api_key(self):
+        """Test initialization fails without API key."""
+        with pytest.raises(ValueError, match="API key is required"):
+            LLMClient(api_key=None)
+
+    def test_get_headers(self):
+        """Test getting headers."""
+        client = LLMClient(api_key="sk-test")
+        headers = client._get_headers()
+        assert headers["Authorization"] == "Bearer sk-test"
+        assert headers["Content-Type"] == "application/json"
+
+    @patch("code_graph_builder.rag.client.requests.post")
+    def test_chat_success(self, mock_post):
+        """Test successful chat completion."""
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            "choices": [
+                {
+                    "message": {"content": "Test response"},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {"prompt_tokens": 10, "completion_tokens": 20},
+            "model": "kimi-k2.5",
+        }
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        client = LLMClient(api_key="sk-test")
+        response = client.chat("Hello")
+
+        assert response.content == "Test response"
+        assert response.model == "kimi-k2.5"
+        mock_post.assert_called_once()
+
+    @patch("code_graph_builder.rag.client.requests.post")
+    def test_chat_with_context(self, mock_post):
+        """Test chat with context."""
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            "choices": [
+                {
+                    "message": {"content": "Response"},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {},
+            "model": "kimi-k2.5",
+        }
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        client = LLMClient(api_key="sk-test")
+        response = client.chat(
+            query="Explain",
+            context="def foo(): pass",
+            system_prompt="You are helpful.",
+        )
+
+        assert response.content == "Response"
+        # Check that context was included in the call
+        call_args = mock_post.call_args
+        json_data = call_args.kwargs["json"]
+        assert any("Context:" in msg["content"] for msg in json_data["messages"])
+
+    @patch("code_graph_builder.rag.client.requests.post")
+    def test_chat_http_error(self, mock_post):
+        """Test chat with HTTP error."""
+        from requests.exceptions import HTTPError
+
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            "error": {"message": "Invalid API key"}
+        }
+        mock_response.raise_for_status.side_effect = HTTPError(
+            "401 Client Error",
+            response=mock_response,
+        )
+        mock_post.return_value = mock_response
+
+        client = LLMClient(api_key="sk-test")
+        with pytest.raises(RuntimeError, match="API request failed"):
+            client.chat("Hello")
+
+    @patch("code_graph_builder.rag.client.requests.post")
+    def test_chat_timeout(self, mock_post):
+        """Test chat with timeout."""
+        from requests.exceptions import Timeout
+
+        mock_post.side_effect = Timeout("Request timed out")
+
+        client = LLMClient(api_key="sk-test", timeout=5)
+        with pytest.raises(RuntimeError, match="timeout"):
+            client.chat("Hello")
+
+    @patch("code_graph_builder.rag.client.requests.post")
+    def test_chat_with_messages(self, mock_post):
+        """Test chat with raw messages."""
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            "choices": [
+                {
+                    "message": {"content": "Response"},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {},
+            "model": "kimi-k2.5",
+        }
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        client = LLMClient(api_key="sk-test")
+        messages = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "Hello"},
+        ]
+        response = client.chat_with_messages(messages)
+
+        assert response.content == "Response"
+        call_args = mock_post.call_args
+        json_data = call_args.kwargs["json"]
+        assert json_data["messages"] == messages
+
+    @patch("code_graph_builder.rag.client.requests.get")
+    def test_health_check_success(self, mock_get):
+        """Test successful health check."""
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_get.return_value = mock_response
+
+        client = LLMClient(api_key="sk-test")
+        assert client.health_check() is True
+
+    @patch("code_graph_builder.rag.client.requests.get")
+    def test_health_check_failure(self, mock_get):
+        """Test failed health check."""
+        mock_get.side_effect = Exception("Connection error")
+
+        client = LLMClient(api_key="sk-test")
+        assert client.health_check() is False
+
+
+class TestCreateLLMClient:
+    """Tests for create_llm_client factory function."""
+
+    def test_create_with_defaults(self):
+        """Test creating client with defaults."""
+        client = create_llm_client(api_key="sk-test")
+        assert isinstance(client, LLMClient)
+        assert client.model == "kimi-k2.5"
+
+    def test_create_with_custom_model(self):
+        """Test creating client with custom model."""
+        client = create_llm_client(
+            api_key="sk-test",
+            model="custom-model",
+            max_tokens=2048,
+        )
+        assert client.model == "custom-model"
+        assert client.max_tokens == 2048

code_graph_builder/rag/tests/test_config.py

@@ -0,0 +1,177 @@
+"""Tests for RAG configuration."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+
+from code_graph_builder.rag.config import (
+    MoonshotConfig,
+    OutputConfig,
+    RAGConfig,
+    RetrievalConfig,
+)
+
+
+class TestMoonshotConfig:
+    """Tests for MoonshotConfig."""
+
+    def test_default_values(self):
+        """Test default configuration values."""
+        config = MoonshotConfig(api_key="sk-test")
+        assert config.model == "kimi-k2.5"
+        assert config.base_url == "https://api.moonshot.cn/v1"
+        assert config.max_tokens == 4096
+        assert config.temperature == 0.7
+        assert config.timeout == 120
+
+    def test_custom_values(self):
+        """Test custom configuration values."""
+        config = MoonshotConfig(
+            api_key="sk-test",
+            model="kimi-k2.5",
+            base_url="https://custom.api.com",
+            max_tokens=2048,
+            temperature=0.5,
+            timeout=60,
+        )
+        assert config.model == "kimi-k2.5"
+        assert config.base_url == "https://custom.api.com"
+        assert config.max_tokens == 2048
+        assert config.temperature == 0.5
+        assert config.timeout == 60
+
+    def test_api_key_from_env(self, monkeypatch):
+        """Test loading API key from environment."""
+        monkeypatch.setenv("MOONSHOT_API_KEY", "sk-from-env")
+        config = MoonshotConfig()
+        assert config.api_key == "sk-from-env"
+
+    def test_validate_missing_api_key(self):
+        """Test validation fails with missing API key."""
+        config = MoonshotConfig(api_key=None)
+        with pytest.raises(ValueError, match="API key is required"):
+            config.validate()
+
+    def test_validate_invalid_api_key_format(self):
+        """Test validation fails with invalid API key format."""
+        config = MoonshotConfig(api_key="invalid-key")
+        with pytest.raises(ValueError, match="start with 'sk-'"):
+            config.validate()
+
+    def test_validate_invalid_temperature(self):
+        """Test validation fails with invalid temperature."""
+        config = MoonshotConfig(api_key="sk-test", temperature=3.0)
+        with pytest.raises(ValueError, match="between 0 and 2"):
+            config.validate()
+
+    def test_to_dict(self):
+        """Test conversion to dictionary."""
+        config = MoonshotConfig(api_key="sk-test")
+        data = config.to_dict()
+        assert data["model"] == "kimi-k2.5"
+        assert data["api_key"] == "sk-test"
+        assert "base_url" in data
+
+
+class TestRetrievalConfig:
+    """Tests for RetrievalConfig."""
+
+    def test_default_values(self):
+        """Test default configuration values."""
+        config = RetrievalConfig()
+        assert config.semantic_top_k == 10
+        assert config.graph_max_depth == 2
+        assert config.include_callers is True
+        assert config.include_callees is True
+        assert config.include_related is True
+        assert config.max_context_tokens == 8000
+        assert config.code_chunk_size == 2000
+
+    def test_custom_values(self):
+        """Test custom configuration values."""
+        config = RetrievalConfig(
+            semantic_top_k=20,
+            graph_max_depth=3,
+            include_callers=False,
+            include_callees=False,
+            include_related=False,
+        )
+        assert config.semantic_top_k == 20
+        assert config.graph_max_depth == 3
+        assert config.include_callers is False
+
+    def test_to_dict(self):
+        """Test conversion to dictionary."""
+        config = RetrievalConfig()
+        data = config.to_dict()
+        assert "semantic_top_k" in data
+        assert "graph_max_depth" in data
+
+
+class TestOutputConfig:
+    """Tests for OutputConfig."""
+
+    def test_default_values(self):
+        """Test default configuration values."""
+        config = OutputConfig()
+        assert config.format == "markdown"
+        assert config.include_source_links is True
+        assert config.include_code_snippets is True
+        assert isinstance(config.output_dir, Path)
+
+    def test_custom_output_dir(self):
+        """Test custom output directory."""
+        config = OutputConfig(output_dir="/custom/path")
+        assert isinstance(config.output_dir, Path)
+        assert str(config.output_dir) == "/custom/path"
+
+    def test_to_dict(self):
+        """Test conversion to dictionary."""
+        config = OutputConfig()
+        data = config.to_dict()
+        assert data["format"] == "markdown"
+        assert "output_dir" in data
+
+
+class TestRAGConfig:
+    """Tests for RAGConfig."""
+
+    def test_default_values(self):
+        """Test default configuration values."""
+        config = RAGConfig(moonshot=MoonshotConfig(api_key="sk-test"))
+        assert isinstance(config.moonshot, MoonshotConfig)
+        assert isinstance(config.retrieval, RetrievalConfig)
+        assert isinstance(config.output, OutputConfig)
+        assert config.verbose is False
+
+    def test_from_env(self, monkeypatch):
+        """Test creating config from environment variables."""
+        monkeypatch.setenv("MOONSHOT_API_KEY", "sk-env-key")
+        monkeypatch.setenv("MOONSHOT_MODEL", "kimi-k2.5")
+        monkeypatch.setenv("RAG_SEMANTIC_TOP_K", "15")
+        monkeypatch.setenv("RAG_OUTPUT_FORMAT", "json")
+        monkeypatch.setenv("RAG_VERBOSE", "true")
+
+        config = RAGConfig.from_env()
+        assert config.moonshot.api_key == "sk-env-key"
+        assert config.moonshot.model == "kimi-k2.5"
+        assert config.retrieval.semantic_top_k == 15
+        assert config.output.format == "json"
+        assert config.verbose is True
+
+    def test_validate(self):
+        """Test configuration validation."""
+        config = RAGConfig(moonshot=MoonshotConfig(api_key="sk-test"))
+        config.validate()  # Should not raise
+
+    def test_to_dict(self):
+        """Test conversion to dictionary."""
+        config = RAGConfig(moonshot=MoonshotConfig(api_key="sk-test"))
+        data = config.to_dict()
+        assert "moonshot" in data
+        assert "retrieval" in data
+        assert "output" in data
+        assert "verbose" in data

code_graph_builder/rag/tests/test_markdown_generator.py

@@ -0,0 +1,240 @@
+"""Tests for markdown generator."""
+
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from code_graph_builder.rag.markdown_generator import (
+    AnalysisResult,
+    MarkdownGenerator,
+    SourceReference,
+    create_source_reference_from_context,
+    format_code_block,
+)
+from code_graph_builder.rag.prompt_templates import CodeContext
+
+
+class TestSourceReference:
+    """Tests for SourceReference."""
+
+    def test_basic_creation(self):
+        """Test basic creation."""
+        ref = SourceReference(
+            name="foo",
+            qualified_name="test.foo",
+            file_path="test.py",
+            line_start=10,
+            line_end=20,
+            entity_type="Function",
+        )
+        assert ref.name == "foo"
+        assert ref.qualified_name == "test.foo"
+
+    def test_format_link(self):
+        """Test link formatting."""
+        ref = SourceReference(
+            name="foo",
+            qualified_name="test.foo",
+            file_path="test.py",
+            line_start=10,
+            line_end=20,
+        )
+        link = ref.format_link()
+        assert "[test.foo]" in link
+        assert "test.py:10-20" in link
+
+    def test_format_link_single_line(self):
+        """Test link formatting for single line."""
+        ref = SourceReference(
+            name="foo",
+            qualified_name="test.foo",
+            file_path="test.py",
+            line_start=10,
+            line_end=10,
+        )
+        link = ref.format_link()
+        assert "test.py:10" in link
+        assert "-10" not in link
+
+    def test_to_dict(self):
+        """Test conversion to dictionary."""
+        ref = SourceReference(
+            name="foo",
+            qualified_name="test.foo",
+            file_path="test.py",
+        )
+        data = ref.to_dict()
+        assert data["name"] == "foo"
+        assert data["qualified_name"] == "test.foo"
+
+
+class TestAnalysisResult:
+    """Tests for AnalysisResult."""
+
+    def test_basic_creation(self):
+        """Test basic creation."""
+        result = AnalysisResult(
+            query="Test query",
+            response="Test response",
+        )
+        assert result.query == "Test query"
+        assert result.response == "Test response"
+        assert result.sources == []
+
+    def test_with_sources(self):
+        """Test creation with sources."""
+        sources = [
+            SourceReference(name="foo", qualified_name="test.foo", file_path="test.py"),
+        ]
+        result = AnalysisResult(
+            query="Test",
+            response="Response",
+            sources=sources,
+            metadata={"key": "value"},
+        )
+        assert len(result.sources) == 1
+        assert result.metadata["key"] == "value"
+
+    def test_to_dict(self):
+        """Test conversion to dictionary."""
+        result = AnalysisResult(query="Test", response="Response")
+        data = result.to_dict()
+        assert data["query"] == "Test"
+        assert data["response"] == "Response"
+        assert "timestamp" in data
+
+
+class TestMarkdownGenerator:
+    """Tests for MarkdownGenerator."""
+
+    def test_default_init(self):
+        """Test default initialization."""
+        gen = MarkdownGenerator()
+        assert gen.include_toc is True
+        assert gen.include_timestamp is True
+        assert gen.include_metadata is True
+
+    def test_custom_init(self):
+        """Test custom initialization."""
+        gen = MarkdownGenerator(
+            include_toc=False,
+            include_timestamp=False,
+            include_metadata=False,
+        )
+        assert gen.include_toc is False
+
+    def test_generate_analysis_doc(self):
+        """Test generating analysis document."""
+        gen = MarkdownGenerator()
+        result = AnalysisResult(
+            query="What does this do?",
+            response="It does something.",
+            sources=[
+                SourceReference(
+                    name="foo",
+                    qualified_name="test.foo",
+                    file_path="test.py",
+                    entity_type="Function",
+                ),
+            ],
+        )
+        doc = gen.generate_analysis_doc("Test Analysis", result)
+        assert "# Test Analysis" in doc
+        assert "## Query" in doc
+        assert "What does this do?" in doc
+        assert "## Analysis" in doc
+        assert "It does something." in doc
+        assert "## Sources" in doc
+        assert "test.foo" in doc
+
+    def test_generate_analysis_doc_no_toc(self):
+        """Test generating document without TOC."""
+        gen = MarkdownGenerator(include_toc=False)
+        result = AnalysisResult(query="Test", response="Response")
+        doc = gen.generate_analysis_doc("Test", result)
+        assert "## Table of Contents" not in doc
+
+    def test_generate_code_documentation(self):
+        """Test generating code documentation."""
+        gen = MarkdownGenerator()
+        ctx = CodeContext(
+            source_code="def foo(): pass",
+            file_path="test.py",
+            qualified_name="test.foo",
+            entity_type="Function",
+        )
+        doc = gen.generate_code_documentation(ctx, "This is a test function.")
+        assert "# test.foo" in doc
+        assert "**Type:** Function" in doc
+        assert "**File:** `test.py`" in doc
+        assert "This is a test function." in doc
+        assert "## Source Code" in doc
+
+    def test_generate_comparison_doc(self):
+        """Test generating comparison document."""
+        gen = MarkdownGenerator()
+        contexts = [
+            CodeContext(source_code="def foo(): pass", qualified_name="foo"),
+            CodeContext(source_code="def bar(): pass", qualified_name="bar"),
+        ]
+        doc = gen.generate_comparison_doc(
+            title="Comparison",
+            query="Compare these",
+            contexts=contexts,
+            analysis="They are similar.",
+        )
+        assert "# Comparison" in doc
+        assert "**Query:** Compare these" in doc
+        assert "## Comparison Analysis" in doc
+        assert "## Compared Entities" in doc
+
+    def test_save_document(self):
+        """Test saving document."""
+        gen = MarkdownGenerator()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = gen.save_document("# Test", f"{tmpdir}/test.md")
+            assert path.exists()
+            assert path.read_text() == "# Test"
+
+    def test_save_document_creates_dirs(self):
+        """Test saving document creates directories."""
+        gen = MarkdownGenerator()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = gen.save_document("# Test", f"{tmpdir}/nested/dir/test.md")
+            assert path.exists()
+
+
+class TestConvenienceFunctions:
+    """Tests for convenience functions."""
+
+    def test_create_source_reference_from_context(self):
+        """Test creating source reference from context."""
+        ctx = CodeContext(
+            source_code="def foo(): pass",
+            file_path="test.py",
+            qualified_name="test.foo",
+            entity_type="Function",
+        )
+        ref = create_source_reference_from_context(ctx)
+        assert isinstance(ref, SourceReference)
+        assert ref.name == "foo"
+        assert ref.qualified_name == "test.foo"
+        assert ref.file_path == "test.py"
+        assert ref.entity_type == "Function"
+
+    def test_format_code_block(self):
+        """Test formatting code block."""
+        code = "def foo(): pass"
+        formatted = format_code_block(code)
+        assert formatted.startswith("```python")
+        assert formatted.endswith("```")
+        assert "def foo(): pass" in formatted
+
+    def test_format_code_block_custom_language(self):
+        """Test formatting code block with custom language."""
+        code = "console.log('test')"
+        formatted = format_code_block(code, language="javascript")
+        assert formatted.startswith("```javascript")