Skip to content

RAG System and Auxiliary Types

RAG System

SourceNode

Bases: BaseModel

Source code in src/fed_rag/types/rag_system.py 
13
14
15
16
17
18
19
class SourceNode(BaseModel):
    score: float
    node: KnowledgeNode

    def __getattr__(self, __name: str) -> Any:
        """Convenient wrapper on getattr of associated node."""
        return getattr(self.node, __name)

__getattr__ 

__getattr__(__name)

Convenient wrapper on getattr of associated node.

Source code in src/fed_rag/types/rag_system.py 
17
18
19
def __getattr__(self, __name: str) -> Any:
    """Convenient wrapper on getattr of associated node."""
    return getattr(self.node, __name)

RAGSystem

Bases: BaseModel

Source code in src/fed_rag/types/rag_system.py 
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
class RAGSystem(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    generator: BaseGenerator
    retriever: BaseRetriever
    knowledge_store: BaseKnowledgeStore
    rag_config: RAGConfig

    def query(self, query: str) -> RAGResponse:
        """Query the RAG system."""
        source_nodes = self.retrieve(query)
        context = self._format_context(source_nodes)
        response = self.generate(query=query, context=context)
        return RAGResponse(source_nodes=source_nodes, response=response)

    def retrieve(self, query: str) -> list[SourceNode]:
        """Retrieve from KnowledgeStore."""
        query_emb: list[float] = self.retriever.encode_query(query).tolist()
        raw_retrieval_result = self.knowledge_store.retrieve(
            query_emb=query_emb, top_k=self.rag_config.top_k
        )
        return [
            SourceNode(score=el[0], node=el[1]) for el in raw_retrieval_result
        ]

    def generate(self, query: str, context: str) -> str:
        """Generate response to query with context."""
        return self.generator.generate(query=query, context=context)  # type: ignore

    def _format_context(self, source_nodes: list[SourceNode]) -> str:
        """Format the context from the source nodes."""
        # TODO: how to format image context
        return self.rag_config.context_separator.join(
            [node.get_content()["text_content"] for node in source_nodes]
        )

query 

query(query)

Query the RAG system.

Source code in src/fed_rag/types/rag_system.py 
42
43
44
45
46
47
def query(self, query: str) -> RAGResponse:
    """Query the RAG system."""
    source_nodes = self.retrieve(query)
    context = self._format_context(source_nodes)
    response = self.generate(query=query, context=context)
    return RAGResponse(source_nodes=source_nodes, response=response)

retrieve 

retrieve(query)

Retrieve from KnowledgeStore.

Source code in src/fed_rag/types/rag_system.py 
49
50
51
52
53
54
55
56
57
def retrieve(self, query: str) -> list[SourceNode]:
    """Retrieve from KnowledgeStore."""
    query_emb: list[float] = self.retriever.encode_query(query).tolist()
    raw_retrieval_result = self.knowledge_store.retrieve(
        query_emb=query_emb, top_k=self.rag_config.top_k
    )
    return [
        SourceNode(score=el[0], node=el[1]) for el in raw_retrieval_result
    ]

generate 

generate(query, context)

Generate response to query with context.

Source code in src/fed_rag/types/rag_system.py 
59
60
61
def generate(self, query: str, context: str) -> str:
    """Generate response to query with context."""
    return self.generator.generate(query=query, context=context)  # type: ignore