Skip to content

Conversation

@gpernelle
Copy link
Contributor

Directly Query for Pydantic Objects with GraphRAG

This PR introduces a powerful new feature that allows you to directly receive a Pydantic object as the result of a GraphRAG query. This significantly streamlines data handling and validation, enabling you to work with structured, type-hinted outputs immediately.

How to Use

By simply passing a Pydantic BaseModel as the response_model argument to the grag.query() method, GraphRAG will automatically parse and validate the query's output into an instance of your specified model.

Here's an example:

from typing import Annotated, List, Optional
from pydantic import BaseModel, Field
from fast_graphrag import GraphRAG

class GraphRAGTheme(BaseModel):
    """A single theme generated by GraphRAG."""
    theme_name: str = Field(description="Name of the theme")
    keywords: List[str] = Field(description="Keywords associated with the theme")

class GraphRAGThemeList(BaseModel):
    """List of themes generated by GraphRAG."""
    themes: List[GraphRAGTheme] = Field(description="List of generated themes")
    city_name: str = Field(description="Name of the city")
    total_themes: int = Field(description="Total number of themes generated")

query = 'your query....'
grag = GraphRAG(....)

# Directly get a Pydantic object from the query result
result: GraphRAGThemeList = grag.query(query, response_model=GraphRAGThemeList)

# You can now access your data with full type hinting and validation
print(result.city_name)
for theme in result.themes:
    print(theme.theme_name)

@liukidar liukidar merged commit cb13227 into circlemind-ai:main Jun 21, 2025
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants