Add models and utils modules
This commit is contained in:
94
src/models/document.py
Normal file
94
src/models/document.py
Normal file
@@ -0,0 +1,94 @@
|
||||
"""Document models for indexed documentation."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class SourceType(str, Enum):
|
||||
"""Enumeration of supported documentation source types."""
|
||||
|
||||
OPENAPI = "openapi"
|
||||
README = "readme"
|
||||
CODE = "code"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""Represents an indexed document chunk."""
|
||||
|
||||
id: str
|
||||
content: str
|
||||
source_type: SourceType
|
||||
title: str
|
||||
file_path: str = ""
|
||||
metadata: dict = field(default_factory=dict)
|
||||
created_at: datetime = field(default_factory=datetime.utcnow)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert document to dictionary for serialization."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"content": self.content,
|
||||
"source_type": self.source_type.value,
|
||||
"title": self.title,
|
||||
"file_path": self.file_path,
|
||||
"metadata": self.metadata,
|
||||
"created_at": self.created_at.isoformat(),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> "Document":
|
||||
"""Create document from dictionary."""
|
||||
return cls(
|
||||
id=data["id"],
|
||||
content=data["content"],
|
||||
source_type=SourceType(data["source_type"]),
|
||||
title=data["title"],
|
||||
file_path=data.get("file_path", ""),
|
||||
metadata=data.get("metadata", {}),
|
||||
created_at=datetime.fromisoformat(data["created_at"]),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""Represents a search result with relevance score."""
|
||||
|
||||
document: Document
|
||||
score: float
|
||||
highlights: list[str] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert search result to dictionary."""
|
||||
return {
|
||||
"id": self.document.id,
|
||||
"content": self.document.content,
|
||||
"source_type": self.document.source_type.value,
|
||||
"title": self.document.title,
|
||||
"file_path": self.document.file_path,
|
||||
"score": self.score,
|
||||
"highlights": self.highlights,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class IndexStats:
|
||||
"""Statistics about the indexed collection."""
|
||||
|
||||
total_documents: int = 0
|
||||
openapi_count: int = 0
|
||||
readme_count: int = 0
|
||||
code_count: int = 0
|
||||
last_indexed: Optional[datetime] = None
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert stats to dictionary."""
|
||||
return {
|
||||
"total_documents": self.total_documents,
|
||||
"openapi_count": self.openapi_count,
|
||||
"readme_count": self.readme_count,
|
||||
"code_count": self.code_count,
|
||||
"last_indexed": self.last_indexed.isoformat() if self.last_indexed else None,
|
||||
}
|
||||
Reference in New Issue
Block a user