docugraph 0.1.0__py3-none-any.whl

docugraph/core/models.py

@@ -0,0 +1,145 @@
+"""Core data models for DocuGraph AI."""
+
+from datetime import UTC, datetime
+from enum import StrEnum
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+def generate_id() -> str:
+    """Generate a unique ID."""
+    return str(uuid4())
+
+
+class ContentType(StrEnum):
+    """Type of document content."""
+
+    MARKDOWN = "markdown"
+    CODE = "code"
+    API_SPEC = "api_spec"
+    TEXT = "text"
+
+
+class EntityType(StrEnum):
+    """Type of knowledge graph entity."""
+
+    FUNCTION = "function"
+    CLASS = "class"
+    MODULE = "module"
+    PACKAGE = "package"
+    ENDPOINT = "endpoint"
+    PARAMETER = "parameter"
+    CONCEPT = "concept"
+    PATTERN = "pattern"
+    DOCUMENTATION = "documentation"
+
+
+class RelationshipType(StrEnum):
+    """Type of relationship between entities."""
+
+    CALLS = "CALLS"
+    IMPORTS = "IMPORTS"
+    INHERITS = "INHERITS"
+    IMPLEMENTS = "IMPLEMENTS"
+    DEPENDS_ON = "DEPENDS_ON"
+    USED_BY = "USED_BY"
+    DESCRIBES = "DESCRIBES"
+    EXAMPLE_OF = "EXAMPLE_OF"
+    RELATED_TO = "RELATED_TO"
+    SUPERSEDES = "SUPERSEDES"
+    DEPRECATES = "DEPRECATES"
+
+
+class Document(BaseModel):
+    """A crawled or indexed document."""
+
+    id: str = Field(default_factory=generate_id)
+    source_url: str | None = None
+    source_path: str | None = None
+    title: str | None = None
+    content: str
+    content_type: ContentType = ContentType.MARKDOWN
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+
+class Chunk(BaseModel):
+    """A chunk of a document for vector storage."""
+
+    id: str = Field(default_factory=generate_id)
+    document_id: str
+    content: str
+    embedding: list[float] | None = None
+    start_char: int = 0
+    end_char: int = 0
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+
+class Entity(BaseModel):
+    """A knowledge graph entity (node)."""
+
+    id: str = Field(default_factory=generate_id)
+    name: str
+    entity_type: EntityType
+    description: str | None = None
+    source_chunks: list[str] = Field(default_factory=list)
+    properties: dict[str, Any] = Field(default_factory=dict)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    valid_from: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    valid_to: datetime | None = None
+
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+
+class Relationship(BaseModel):
+    """A knowledge graph relationship (edge)."""
+
+    id: str = Field(default_factory=generate_id)
+    source_id: str
+    target_id: str
+    relationship_type: RelationshipType
+    properties: dict[str, Any] = Field(default_factory=dict)
+    weight: float = 1.0
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+
+class MemoryEntry(BaseModel):
+    """An entry in the agent's memory store."""
+
+    id: str = Field(default_factory=generate_id)
+    repository: str
+    branch: str = "main"
+    key: str
+    value: Any
+    context: dict[str, Any] = Field(default_factory=dict)
+    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+    accessed_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
+
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+
+class SearchResult(BaseModel):
+    """A search result with relevance score."""
+
+    # Class-based `Config` is deprecated in Pydantic v2 and warns on import.
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    chunk: Chunk
+    score: float
+    source_document: Document | None = None
+    highlights: list[str] = Field(default_factory=list)

docugraph/ingestion/__init__.py

@@ -0,0 +1,64 @@
+"""Ingestion pipeline: crawling, file processing, and chunking."""
+
+from docugraph.ingestion.api_docs import (
+    APIDocConfig,
+    APIDocParser,
+    APIDocType,
+    APIEndpoint,
+    APIParameter,
+    APISchema,
+    GraphQLType,
+    ParsedAPISpec,
+    get_api_doc_parser,
+)
+from docugraph.ingestion.chunker import (
+    Chunker,
+    ChunkerConfig,
+    ChunkingStrategy,
+    get_chunker,
+)
+from docugraph.ingestion.crawler import DocCrawler
+from docugraph.ingestion.git_indexer import (
+    GitIndexer,
+    GitIndexerConfig,
+    RepoInfo,
+    get_git_indexer,
+)
+from docugraph.ingestion.local_files import (
+    FileType,
+    IndexedFile,
+    LocalFileConfig,
+    LocalFileIndexer,
+    get_local_file_indexer,
+)
+
+__all__ = [
+    # Crawler
+    "DocCrawler",
+    # Chunker
+    "Chunker",
+    "ChunkerConfig",
+    "ChunkingStrategy",
+    "get_chunker",
+    # Local Files
+    "FileType",
+    "IndexedFile",
+    "LocalFileConfig",
+    "LocalFileIndexer",
+    "get_local_file_indexer",
+    # Git Indexer
+    "GitIndexer",
+    "GitIndexerConfig",
+    "RepoInfo",
+    "get_git_indexer",
+    # API Docs
+    "APIDocType",
+    "APIParameter",
+    "APIEndpoint",
+    "APISchema",
+    "GraphQLType",
+    "ParsedAPISpec",
+    "APIDocConfig",
+    "APIDocParser",
+    "get_api_doc_parser",
+]