from upsonic import Agent, Task, KnowledgeBase
from upsonic.loaders.json import JSONLoader
from upsonic.loaders.config import JSONLoaderConfig
from upsonic.embeddings import OpenAIEmbedding, OpenAIEmbeddingConfig
from upsonic.text_splitter.json_chunker import JSONChunker, JSONChunkingConfig
from upsonic.vectordb import ChromaProvider, ChromaConfig, ConnectionConfig, Mode
# Configure splitter
splitter_config = JSONChunkingConfig(
chunk_size=512,
chunk_overlap=50,
convert_lists_to_dicts=True,
max_depth=50
)
splitter = JSONChunker(splitter_config)
# Setup KnowledgeBase
loader = JSONLoader(JSONLoaderConfig())
embedding = OpenAIEmbedding(OpenAIEmbeddingConfig())
vectordb = ChromaProvider(ChromaConfig(
collection_name="json_docs",
vector_size=1536,
connection=ConnectionConfig(mode=Mode.IN_MEMORY)
))
kb = KnowledgeBase(
sources=["data.json"],
embedding_provider=embedding,
vectordb=vectordb,
loaders=[loader],
splitters=[splitter]
)
# Query with Agent
agent = Agent("openai/gpt-4o")
task = Task("Find all user records", context=[kb])
result = agent.do(task)
print(result)
