1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
"""Abstractions for search functionality."""
import uuid
from typing import Any, Dict, List, Optional, Tuple
from pydantic import BaseModel, Field
from .llm import GenerationConfig
class VectorSearchRequest(BaseModel):
"""Request for a search operation."""
query: str
limit: int
filters: Optional[dict[str, Any]] = None
class VectorSearchResult(BaseModel):
"""Result of a search operation."""
id: uuid.UUID
score: float
metadata: dict[str, Any]
def __str__(self) -> str:
return f"VectorSearchResult(id={self.id}, score={self.score}, metadata={self.metadata})"
def __repr__(self) -> str:
return f"VectorSearchResult(id={self.id}, score={self.score}, metadata={self.metadata})"
def dict(self) -> dict:
return {
"id": self.id,
"score": self.score,
"metadata": self.metadata,
}
class KGSearchRequest(BaseModel):
"""Request for a knowledge graph search operation."""
query: str
# [query, ...]
KGSearchResult = List[Tuple[str, List[Dict[str, Any]]]]
class AggregateSearchResult(BaseModel):
"""Result of an aggregate search operation."""
vector_search_results: Optional[List[VectorSearchResult]]
kg_search_results: Optional[KGSearchResult] = None
def __str__(self) -> str:
return f"AggregateSearchResult(vector_search_results={self.vector_search_results}, kg_search_results={self.kg_search_results})"
def __repr__(self) -> str:
return f"AggregateSearchResult(vector_search_results={self.vector_search_results}, kg_search_results={self.kg_search_results})"
def dict(self) -> dict:
return {
"vector_search_results": (
[result.dict() for result in self.vector_search_results]
if self.vector_search_results
else []
),
"kg_search_results": self.kg_search_results or [],
}
class VectorSearchSettings(BaseModel):
use_vector_search: bool = True
search_filters: dict[str, Any] = Field(default_factory=dict)
search_limit: int = 10
do_hybrid_search: bool = False
class KGSearchSettings(BaseModel):
use_kg_search: bool = False
agent_generation_config: Optional[GenerationConfig] = Field(
default_factory=GenerationConfig
)
|
