from upsonic import Agent, Task, KnowledgeBase
from upsonic.embeddings import OpenAIEmbedding, OpenAIEmbeddingConfig
from upsonic.vectordb import ChromaProvider, ChromaConfig, ConnectionConfig, Mode
embedding = OpenAIEmbedding(OpenAIEmbeddingConfig())
# Both KBs point to the same collection
def make_vectordb():
return ChromaProvider(ChromaConfig(
collection_name="company_docs",
vector_size=1536,
connection=ConnectionConfig(mode=Mode.EMBEDDED, db_path="./chroma_db"),
))
# KB for HR policies
hr_kb = KnowledgeBase(
sources=["hr_handbook.pdf"],
vectordb=make_vectordb(),
embedding_provider=embedding,
name="hr_policies",
isolate_search=True, # default — only returns HR docs
)
# KB for engineering docs
eng_kb = KnowledgeBase(
sources=["engineering_wiki/"],
vectordb=make_vectordb(),
embedding_provider=embedding,
name="engineering_docs",
isolate_search=True, # default — only returns engineering docs
)
agent = Agent("anthropic/claude-sonnet-4-5")
# This query only searches HR documents
hr_task = Task(
description="What is the PTO policy?",
context=[hr_kb],
)
print(agent.do(hr_task))
# This query only searches engineering documents
eng_task = Task(
description="How do we deploy to production?",
context=[eng_kb],
)
print(agent.do(eng_task))