Coverage for src / lexigram / contracts / ai / index.py: 0%

1"""Index and query engine protocols for RAG.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import AsyncIterator 

6from dataclasses import dataclass, field 

7from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable 

8 

9if TYPE_CHECKING: 

10 from lexigram.contracts.ai.vector import Document 

11 from lexigram.contracts.core.result import Result 

12 

13from lexigram.contracts.ai.exceptions import RAGError 

14 

15 

16class IndexError(RAGError): # noqa: A001 

17 """Raised when index operations fail in an expected, recoverable way.""" 

18 

19 _code = "LEX_ERR_IDX_001" 

20 

21 def __init__(self, message: str = "Index error", **kwargs: Any) -> None: 

22 super().__init__(message, **kwargs) 

23 

24 

25class QueryEngineError(RAGError): 

26 """Raised when query engine operations fail in an expected, recoverable way.""" 

27 

28 _code = "LEX_ERR_QE_001" 

29 

30 def __init__(self, message: str = "Query engine error", **kwargs: Any) -> None: 

31 super().__init__(message, **kwargs) 

32 

33 

34@dataclass(frozen=True) 

35class Citation: 

36 """A citation linking an answer fragment to a source node. 

37 

38 Attributes: 

39 node_id: Unique identifier of the source node. 

40 text: The cited text content. 

41 score: Relevance score of the citation. 

42 """ 

43 

44 node_id: str 

45 text: str 

46 score: float 

47 

48 

49@dataclass(frozen=True) 

50class QueryEngineResponse: 

51 """Response from a query engine. 

52 

53 Attributes: 

54 answer: The generated answer text. 

55 source_nodes: List of source nodes used in generation. 

56 citations: List of citations linking answer to sources. 

57 tokens: Total tokens used (if available, otherwise 0). 

58 cost: Total cost in USD (if available, otherwise 0.0). 

59 """ 

60 

61 answer: str 

62 source_nodes: list[Any] = field(default_factory=list) 

63 citations: list[Citation] = field(default_factory=list) 

64 tokens: int = 0 

65 cost: float = 0.0 

66 

67 

68@runtime_checkable 

69class IndexProtocol(Protocol): 

70 """Protocol for document indices. 

71 

72 Implementations store embedded documents and provide search capabilities. 

73 """ 

74 

75 async def insert(self, documents: list[Document]) -> Result[list[str], IndexError]: 

76 """Insert documents into the index. 

77 

78 Args: 

79 documents: List of documents to insert. 

80 

81 Returns: 

82 Ok(list of document IDs) on success. 

83 Err(IndexError) on failure. 

84 """ 

85 ... 

86 

87 async def delete(self, ids: list[str]) -> Result[int, IndexError]: 

88 """Delete documents by ID. 

89 

90 Args: 

91 ids: List of document IDs to delete. 

92 

93 Returns: 

94 Ok(count of deleted documents) on success. 

95 Err(IndexError) on failure. 

96 """ 

97 ... 

98 

99 def as_retriever(self, **kwargs: Any) -> Any: 

100 """Convert this index to a retriever. 

101 

102 Args: 

103 **kwargs: Retriever-specific parameters (top_k, filters, etc.). 

104 

105 Returns: 

106 A retriever instance. 

107 """ 

108 ... 

109 

110 

111@runtime_checkable 

112class QueryEngineProtocol(Protocol): 

113 """Protocol for query engines. 

114 

115 Implementations process user queries and return answers with sources. 

116 """ 

117 

118 async def query( 

119 self, query: str, **kwargs: Any 

120 ) -> Result[QueryEngineResponse, QueryEngineError]: 

121 """Process a query and return an answer with sources. 

122 

123 Args: 

124 query: The user's query string. 

125 **kwargs: Query-specific parameters. 

126 

127 Returns: 

128 Ok(QueryEngineResponse) on success. 

129 Err(QueryEngineError) on failure. 

130 """ 

131 ... 

132 

133 async def astream_query( 

134 self, query: str, **kwargs: Any 

135 ) -> AsyncIterator[Result[QueryEngineResponse, QueryEngineError]]: 

136 """Stream query results. 

137 

138 Args: 

139 query: The user's query string. 

140 **kwargs: Query-specific parameters. 

141 

142 Yields: 

143 Result containing partial QueryEngineResponse or error. 

144 """ 

145 ... 

146 

147 

148__all__ = [ 

149 "Citation", 

150 "IndexError", 

151 "IndexProtocol", 

152 "QueryEngineError", 

153 "QueryEngineProtocol", 

154 "QueryEngineResponse", 

155]