railtracks.prebuilt

 8def rag_node(
 9    documents: List[str],
10    embed_model="text-embedding-3-small",
11    token_count_model="gpt-4o",
12    chunk_size=1000,
13    chunk_overlap=200,
14):
15    """
16    Creates a rag node that allows you to vector the search the provided documents.
17
18    Args:
19        documents (list): List of documents to process. Each document is a string of raw text.
20        embed_model (str): Model name for embedding service.
21        token_count_model (str): Model name for token counting.
22        chunk_size (int): Size of each text chunk.
23        chunk_overlap (int): Overlap between chunks.
24
25    Returns:
26        A node to be invoked upon request.
27
28    """
29
30    rag_core = RAG(
31        docs=documents,
32        config=RAGConfig(
33            embedding={"model": embed_model},
34            store={},
35            chunking={
36                "chunk_size": chunk_size,
37                "chunk_overlap": chunk_overlap,
38                "model": token_count_model,
39            },
40        ),
41    )
42    rag_core.embed_all()
43
44    def query(query: str, top_k: int = 1) -> SearchResult:
45        result = rag_core.search(query, top_k=top_k)
46        return result
47
48    return rt.function_node(query)