Skip to content

Commit

Permalink
fix: Apply small fixes to docs logic
Browse files Browse the repository at this point in the history
  • Loading branch information
whiterabbit1983 committed Dec 27, 2024
1 parent fece441 commit ea37837
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 10 deletions.
16 changes: 8 additions & 8 deletions agents-api/agents_api/routers/docs/search_docs.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ async def get_search_fn_and_params(
case TextOnlyDocSearchRequest(
text=query, limit=k, metadata_filter=metadata_filter
):
search_fn = await search_docs_by_text
search_fn = search_docs_by_text
params = dict(
query=query,
k=k,
Expand All @@ -44,7 +44,7 @@ async def get_search_fn_and_params(
confidence=confidence,
metadata_filter=metadata_filter,
):
search_fn = await search_docs_by_embedding
search_fn = search_docs_by_embedding
params = dict(
query_embedding=query_embedding,
k=k * 3 if search_params.mmr_strength > 0 else k,
Expand All @@ -60,12 +60,12 @@ async def get_search_fn_and_params(
alpha=alpha,
metadata_filter=metadata_filter,
):
search_fn = await search_docs_hybrid
search_fn = search_docs_hybrid
params = dict(
query=query,
query_embedding=query_embedding,
text_query=query,
embedding=query_embedding,
k=k * 3 if search_params.mmr_strength > 0 else k,
embed_search_options=dict(confidence=confidence),
confidence=confidence,
alpha=alpha,
metadata_filter=metadata_filter,
)
Expand Down Expand Up @@ -97,7 +97,7 @@ async def search_user_docs(
search_fn, params = await get_search_fn_and_params(search_params)

start = time.time()
docs: list[DocReference] = search_fn(
docs: list[DocReference] = await search_fn(
developer_id=x_developer_id,
owners=[("user", user_id)],
**params,
Expand Down Expand Up @@ -148,7 +148,7 @@ async def search_agent_docs(
search_fn, params = await get_search_fn_and_params(search_params)

start = time.time()
docs: list[DocReference] = search_fn(
docs: list[DocReference] = await search_fn(
developer_id=x_developer_id,
owners=[("agent", agent_id)],
**params,
Expand Down
4 changes: 3 additions & 1 deletion agents-api/tests/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from agents_api.queries.developers.create_developer import create_developer
from agents_api.queries.developers.get_developer import get_developer
from agents_api.queries.docs.create_doc import create_doc
from agents_api.queries.docs.get_doc import get_doc
from agents_api.queries.executions.create_execution import create_execution
from agents_api.queries.executions.create_execution_transition import (
create_execution_transition,
Expand Down Expand Up @@ -135,7 +136,7 @@ async def test_file(dsn=pg_dsn, developer=test_developer, user=test_user):
@fixture(scope="test")
async def test_doc(dsn=pg_dsn, developer=test_developer, agent=test_agent):
pool = await create_db_pool(dsn=dsn)
doc = await create_doc(
resp = await create_doc(
developer_id=developer.id,
data=CreateDocRequest(
title="Hello",
Expand All @@ -147,6 +148,7 @@ async def test_doc(dsn=pg_dsn, developer=test_developer, agent=test_agent):
owner_id=agent.id,
connection_pool=pool,
)
doc = await get_doc(developer_id=developer.id, doc_id=resp.id, connection_pool=pool)
return doc


Expand Down
1 change: 0 additions & 1 deletion agents-api/tests/test_chat_routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,5 @@ async def _(
connection_pool=pool,
)

print("-->", type(context), context)
assert isinstance(context, ChatContext)
assert len(context.toolsets) > 0

0 comments on commit ea37837

Please sign in to comment.