from upsonic import Agent, Task, KnowledgeBase
from upsonic.loaders import JSONLoader, JSONLoaderConfig
from upsonic.embeddings import OpenAIEmbedding, OpenAIEmbeddingConfig
from upsonic.text_splitter import RecursiveChunker, RecursiveChunkingConfig
from upsonic.vectordb import ChromaProvider, ChromaConfig, ConnectionConfig, Mode
# Configure loader for multi-document extraction
loader_config = JSONLoaderConfig(
mode="multi",
record_selector=".articles[]",
content_mapper=".title + ' ' + .body",
metadata_mapper={"author": ".author", "date": ".published"}
)
loader = JSONLoader(loader_config)
# Setup KnowledgeBase
embedding = OpenAIEmbedding(OpenAIEmbeddingConfig())
chunker = RecursiveChunker(RecursiveChunkingConfig())
vectordb = ChromaProvider(ChromaConfig(
collection_name="json_data",
vector_size=1536,
connection=ConnectionConfig(mode=Mode.IN_MEMORY)
))
kb = KnowledgeBase(
sources=["articles.json"],
embedding_provider=embedding,
vectordb=vectordb,
loaders=[loader],
splitters=[chunker]
)
# Query with Agent
agent = Agent("openai/gpt-4o")
task = Task("Find articles about AI", context=[kb])
result = agent.do(task)
print(result)