mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-21 19:02:45 -05:00
chore(lint): enable perflint ruff rules (#14025)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
1 parent
5ae465fa19
commit
e746c441bc
98 files changed
+595
-619
No files matched your search
@@ -131,10 +131,11 @@ def filter_tenants_by_range(
|
||||
filtered_tenant_schemas = sorted_tenant_schemas[start_idx:end_idx]
|
||||
|
||||
# Combine with non-tenant schemas and preserve original order
|
||||
filtered_tenants = []
|
||||
for tenant_id in tenant_ids:
|
||||
if tenant_id in filtered_tenant_schemas or tenant_id in non_tenant_schemas:
|
||||
filtered_tenants.append(tenant_id)
|
||||
filtered_tenants = [
|
||||
tenant_id
|
||||
for tenant_id in tenant_ids
|
||||
if tenant_id in filtered_tenant_schemas or tenant_id in non_tenant_schemas
|
||||
]
|
||||
|
||||
return filtered_tenants
|
||||
|
||||
|
||||
@@ -100,9 +100,7 @@ def get_google_drive_documents_from_database() -> list[dict]:
|
||||
""")
|
||||
)
|
||||
|
||||
documents = []
|
||||
for row in result:
|
||||
documents.append({"document_id": row.id})
|
||||
documents = [{"document_id": row.id} for row in result]
|
||||
|
||||
return documents
|
||||
|
||||
|
||||
@@ -103,9 +103,9 @@ def upgrade() -> None:
|
||||
"is_cloud": is_cloud,
|
||||
}
|
||||
|
||||
for key, value in config.items():
|
||||
if key not in ["wiki_page_url"]:
|
||||
new_config[key] = value
|
||||
new_config.update(
|
||||
{key: value for key, value in config.items() if key != "wiki_page_url"}
|
||||
)
|
||||
|
||||
op.execute(
|
||||
connector.update()
|
||||
|
||||
@@ -220,9 +220,11 @@ def check_for_doc_permissions_sync(self: Task, *, tenant_id: str) -> bool | None
|
||||
with get_session_with_current_tenant() as db_session:
|
||||
cc_pairs = get_all_auto_sync_cc_pairs(db_session)
|
||||
|
||||
for cc_pair in cc_pairs:
|
||||
if _is_external_doc_permissions_sync_due(cc_pair):
|
||||
cc_pair_ids_to_sync.append(cc_pair.id)
|
||||
cc_pair_ids_to_sync.extend(
|
||||
cc_pair.id
|
||||
for cc_pair in cc_pairs
|
||||
if _is_external_doc_permissions_sync_due(cc_pair)
|
||||
)
|
||||
|
||||
# Tenant-work-gating hook: refresh this tenant's active-set membership
|
||||
# whenever doc-permission sync has any due cc_pairs to dispatch.
|
||||
|
||||
@@ -206,9 +206,11 @@ def check_for_external_group_sync(self: Task, *, tenant_id: str) -> bool | None:
|
||||
if cc_pair.id != cc_pair_to_remove.id
|
||||
]
|
||||
|
||||
for cc_pair in cc_pairs:
|
||||
if _is_external_group_sync_due(cc_pair):
|
||||
cc_pair_ids_to_sync.append(cc_pair.id)
|
||||
cc_pair_ids_to_sync.extend(
|
||||
cc_pair.id
|
||||
for cc_pair in cc_pairs
|
||||
if _is_external_group_sync_due(cc_pair)
|
||||
)
|
||||
|
||||
# Tenant-work-gating hook: refresh this tenant's active-set membership
|
||||
# whenever external-group sync has any due cc_pairs to dispatch.
|
||||
|
||||
@@ -179,7 +179,6 @@ def fetch_repository_team_slugs(
|
||||
) -> List[str]:
|
||||
"""Fetch team slugs with access to the repository."""
|
||||
logger.info("Fetching team slugs for repository %s", repo.full_name)
|
||||
teams_data: List[str] = []
|
||||
|
||||
team_objs: PaginatedList[Team] | list[Team] = (
|
||||
_run_with_retry(
|
||||
@@ -190,8 +189,7 @@ def fetch_repository_team_slugs(
|
||||
or []
|
||||
)
|
||||
|
||||
for team in team_objs:
|
||||
teams_data.append(team.slug)
|
||||
teams_data: List[str] = [team.slug for team in team_objs]
|
||||
|
||||
logger.info(
|
||||
"Fetched %s team slugs for repository %s", len(teams_data), repo.full_name
|
||||
|
||||
@@ -81,11 +81,12 @@ def _post_query_chunk_censoring(
|
||||
final_chunk_dict[censored_chunk.unique_id] = censored_chunk
|
||||
|
||||
# IMPORTANT: make sure to retain the same ordering as the original `chunks` passed in
|
||||
final_chunk_list: list[InferenceChunk] = []
|
||||
for chunk in chunks:
|
||||
# only if the chunk is in the final censored chunks, add it to the final list
|
||||
# if it is missing, that means it was intentionally left out
|
||||
if chunk.unique_id in final_chunk_dict:
|
||||
final_chunk_list.append(final_chunk_dict[chunk.unique_id])
|
||||
# only if the chunk is in the final censored chunks, add it to the final list
|
||||
# if it is missing, that means it was intentionally left out
|
||||
final_chunk_list: list[InferenceChunk] = [
|
||||
final_chunk_dict[chunk.unique_id]
|
||||
for chunk in chunks
|
||||
if chunk.unique_id in final_chunk_dict
|
||||
]
|
||||
|
||||
return final_chunk_list
|
||||
@@ -893,9 +893,10 @@ def get_sharepoint_external_groups(
|
||||
return external_user_groups
|
||||
|
||||
already_resolved = set(groups_and_members.groups_to_emails.keys())
|
||||
for group in _enumerate_ad_groups_paginated(
|
||||
get_access_token, already_resolved, graph_api_base
|
||||
):
|
||||
external_user_groups.append(group)
|
||||
external_user_groups.extend(
|
||||
_enumerate_ad_groups_paginated(
|
||||
get_access_token, already_resolved, graph_api_base
|
||||
)
|
||||
)
|
||||
|
||||
return external_user_groups
|
||||
@@ -267,15 +267,14 @@ def get_assistant_stats(
|
||||
all_dates = set(daily_messages_map.keys()) | set(daily_unique_users_map.keys())
|
||||
|
||||
# Merge both sets of metrics by date
|
||||
daily_results: list[AssistantDailyUsageResponse] = []
|
||||
for date in sorted(all_dates):
|
||||
daily_results.append(
|
||||
AssistantDailyUsageResponse(
|
||||
date=date,
|
||||
total_messages=daily_messages_map.get(date, 0),
|
||||
total_unique_users=daily_unique_users_map.get(date, 0),
|
||||
)
|
||||
daily_results: list[AssistantDailyUsageResponse] = [
|
||||
AssistantDailyUsageResponse(
|
||||
date=date,
|
||||
total_messages=daily_messages_map.get(date, 0),
|
||||
total_unique_users=daily_unique_users_map.get(date, 0),
|
||||
)
|
||||
for date in sorted(all_dates)
|
||||
]
|
||||
|
||||
# Now pull a single total distinct user count across the entire time range
|
||||
total_msgs = sum(d.total_messages for d in daily_results)
|
||||
|
||||
@@ -171,14 +171,13 @@ class QuestionAnswerPairSnapshot(BaseModel):
|
||||
cls,
|
||||
chat_session_snapshot: ChatSessionSnapshot,
|
||||
) -> list["QuestionAnswerPairSnapshot"]:
|
||||
message_pairs: list[tuple[MessageSnapshot, MessageSnapshot]] = []
|
||||
for ind in range(1, len(chat_session_snapshot.messages), 2):
|
||||
message_pairs.append(
|
||||
(
|
||||
chat_session_snapshot.messages[ind - 1],
|
||||
chat_session_snapshot.messages[ind],
|
||||
)
|
||||
message_pairs: list[tuple[MessageSnapshot, MessageSnapshot]] = [
|
||||
(
|
||||
chat_session_snapshot.messages[ind - 1],
|
||||
chat_session_snapshot.messages[ind],
|
||||
)
|
||||
for ind in range(1, len(chat_session_snapshot.messages), 2)
|
||||
]
|
||||
|
||||
return [
|
||||
cls(
|
||||
|
||||
@@ -214,7 +214,7 @@ def celery_inspect_get_reserved(worker_names: list[str], app: Celery) -> set[str
|
||||
# get the list of reserved tasks
|
||||
reserved_tasks: dict[str, list] | None = inspect.reserved()
|
||||
if reserved_tasks:
|
||||
for _, task_list in reserved_tasks.items():
|
||||
for task_list in reserved_tasks.values():
|
||||
for task in task_list:
|
||||
reserved_task_ids.add(task["id"])
|
||||
|
||||
@@ -235,7 +235,7 @@ def celery_inspect_get_active(worker_names: list[str], app: Celery) -> set[str]:
|
||||
# get the list of reserved tasks
|
||||
active_tasks: dict[str, list] | None = inspect.active()
|
||||
if active_tasks:
|
||||
for _, task_list in active_tasks.items():
|
||||
for task_list in active_tasks.values():
|
||||
for task in task_list:
|
||||
active_task_ids.add(task["id"])
|
||||
|
||||
|
||||
@@ -238,8 +238,7 @@ def check_for_pruning(self: Task, *, tenant_id: str) -> bool | None:
|
||||
cc_pair_ids: list[int] = []
|
||||
with get_session_with_current_tenant() as db_session:
|
||||
cc_pairs = get_connector_credential_pairs(db_session)
|
||||
for cc_pair_entry in cc_pairs:
|
||||
cc_pair_ids.append(cc_pair_entry.id)
|
||||
cc_pair_ids.extend(cc_pair_entry.id for cc_pair_entry in cc_pairs)
|
||||
|
||||
prune_dispatched = False
|
||||
for cc_pair_id in cc_pair_ids:
|
||||
|
||||
@@ -159,8 +159,7 @@ def check_for_vespa_sync_task(self: Task, *, tenant_id: str) -> bool | None:
|
||||
db_session=db_session, only_up_to_date=False
|
||||
)
|
||||
|
||||
for usergroup in user_groups:
|
||||
usergroup_ids.append(usergroup.id)
|
||||
usergroup_ids.extend(usergroup.id for usergroup in user_groups)
|
||||
|
||||
for usergroup_id in usergroup_ids:
|
||||
lock_beat.reacquire()
|
||||
|
||||
@@ -512,7 +512,6 @@ def construct_message_history(
|
||||
# Track dropped file messages so we can provide their metadata to the
|
||||
# FileReaderTool instead.
|
||||
truncated_history_before: list[ChatMessageSimple] = []
|
||||
dropped_file_ids: list[str] = []
|
||||
current_token_count = 0
|
||||
|
||||
for msg in reversed(history_before_last_user):
|
||||
@@ -531,9 +530,11 @@ def construct_message_history(
|
||||
# recent messages, so the dropped ones are at the start of the original
|
||||
# list up to (len(history) - len(kept)).
|
||||
num_kept = len(truncated_history_before)
|
||||
for msg in history_before_last_user[: len(history_before_last_user) - num_kept]:
|
||||
if msg.file_id is not None:
|
||||
dropped_file_ids.append(msg.file_id)
|
||||
dropped_file_ids: list[str] = [
|
||||
msg.file_id
|
||||
for msg in history_before_last_user[: len(history_before_last_user) - num_kept]
|
||||
if msg.file_id is not None
|
||||
]
|
||||
|
||||
# Also treat "orphaned" metadata entries as dropped -- these are files
|
||||
# from messages removed by summary truncation (before convert_chat_history
|
||||
@@ -663,10 +664,10 @@ def _create_file_tool_metadata_message(
|
||||
"read sections of any file. You MUST pass the file_id UUID (not the "
|
||||
"filename) to read_file:"
|
||||
]
|
||||
for meta in file_metadata:
|
||||
lines.append(
|
||||
f'- file_id="{meta.file_id}" filename="{meta.filename}" (~{meta.approx_char_count:,} chars)'
|
||||
)
|
||||
lines.extend(
|
||||
f'- file_id="{meta.file_id}" filename="{meta.filename}" (~{meta.approx_char_count:,} chars)'
|
||||
for meta in file_metadata
|
||||
)
|
||||
|
||||
message_content = "\n".join(lines)
|
||||
return ChatMessageSimple(
|
||||
|
||||
@@ -212,8 +212,7 @@ def _get_forums(
|
||||
pages_fetched += len(contents)
|
||||
logger.debug("Fetched %s forums", pages_fetched)
|
||||
|
||||
for page in contents:
|
||||
pages_to_return.append(page)
|
||||
pages_to_return.extend(contents)
|
||||
|
||||
if pages_fetched >= total_records:
|
||||
break
|
||||
|
||||
@@ -1021,33 +1021,31 @@ class CanvasConnector(
|
||||
start: SecondsSinceUnixEpoch | None,
|
||||
end: SecondsSinceUnixEpoch | None,
|
||||
) -> list[SlimDocument]:
|
||||
slim_docs: list[SlimDocument] = []
|
||||
|
||||
for page in self._list_pages(course_id):
|
||||
slim_docs.append(
|
||||
SlimDocument(
|
||||
id=f"canvas-page-{page.course_id}-{page.page_id}",
|
||||
external_access=self._get_item_permissions(course_id, page)
|
||||
or ExternalAccess.empty(),
|
||||
)
|
||||
slim_docs: list[SlimDocument] = [
|
||||
SlimDocument(
|
||||
id=f"canvas-page-{page.course_id}-{page.page_id}",
|
||||
external_access=self._get_item_permissions(course_id, page)
|
||||
or ExternalAccess.empty(),
|
||||
)
|
||||
for page in self._list_pages(course_id)
|
||||
]
|
||||
|
||||
for assignment in self._list_assignments(course_id):
|
||||
slim_docs.append(
|
||||
SlimDocument(
|
||||
id=f"canvas-assignment-{assignment.course_id}-{assignment.id}",
|
||||
external_access=self._get_item_permissions(course_id, assignment)
|
||||
or ExternalAccess.empty(),
|
||||
)
|
||||
slim_docs.extend(
|
||||
SlimDocument(
|
||||
id=f"canvas-assignment-{assignment.course_id}-{assignment.id}",
|
||||
external_access=self._get_item_permissions(course_id, assignment)
|
||||
or ExternalAccess.empty(),
|
||||
)
|
||||
for assignment in self._list_assignments(course_id)
|
||||
)
|
||||
|
||||
for announcement in self._list_announcements(course_id, start, end):
|
||||
slim_docs.append(
|
||||
SlimDocument(
|
||||
id=f"canvas-announcement-{announcement.course_id}-{announcement.id}",
|
||||
external_access=self._get_item_permissions(course_id, announcement)
|
||||
or ExternalAccess.empty(),
|
||||
)
|
||||
slim_docs.extend(
|
||||
SlimDocument(
|
||||
id=f"canvas-announcement-{announcement.course_id}-{announcement.id}",
|
||||
external_access=self._get_item_permissions(course_id, announcement)
|
||||
or ExternalAccess.empty(),
|
||||
)
|
||||
for announcement in self._list_announcements(course_id, start, end)
|
||||
)
|
||||
|
||||
return slim_docs
|
||||
@@ -231,9 +231,7 @@ class CodaConnector(LoadConnector, PollConnector):
|
||||
else:
|
||||
raise
|
||||
|
||||
values = {}
|
||||
for col_name, col_value in response.get("values", {}).items():
|
||||
values[col_name] = col_value
|
||||
values = dict(response.get("values", {}))
|
||||
|
||||
return CodaRow(
|
||||
id=response["id"],
|
||||
@@ -408,17 +406,17 @@ class CodaConnector(LoadConnector, PollConnector):
|
||||
raise
|
||||
|
||||
items = response.get("items", [])
|
||||
for item in items:
|
||||
tables.append(
|
||||
CodaTable(
|
||||
id=item["id"],
|
||||
browser_link=item["browserLink"],
|
||||
name=item["name"],
|
||||
created_at=item["createdAt"],
|
||||
updated_at=item["updatedAt"],
|
||||
doc_id=doc_id,
|
||||
)
|
||||
tables.extend(
|
||||
CodaTable(
|
||||
id=item["id"],
|
||||
browser_link=item["browserLink"],
|
||||
name=item["name"],
|
||||
created_at=item["createdAt"],
|
||||
updated_at=item["updatedAt"],
|
||||
doc_id=doc_id,
|
||||
)
|
||||
for item in items
|
||||
)
|
||||
|
||||
next_page_token = response.get("nextPageToken")
|
||||
if not next_page_token:
|
||||
@@ -453,9 +451,7 @@ class CodaConnector(LoadConnector, PollConnector):
|
||||
|
||||
items = response.get("items", [])
|
||||
for item in items:
|
||||
values = {}
|
||||
for col_name, col_value in item.get("values", {}).items():
|
||||
values[col_name] = col_value
|
||||
values = dict(item.get("values", {}))
|
||||
|
||||
rows.append(
|
||||
CodaRow(
|
||||
|
||||
@@ -570,8 +570,10 @@ class ConfluenceConnector(
|
||||
# Extract labels
|
||||
labels = []
|
||||
if "metadata" in page and "labels" in page["metadata"]:
|
||||
for label in page["metadata"]["labels"].get("results", []):
|
||||
labels.append(label.get("name", ""))
|
||||
labels.extend(
|
||||
label.get("name", "")
|
||||
for label in page["metadata"]["labels"].get("results", [])
|
||||
)
|
||||
if labels:
|
||||
metadata["labels"] = labels
|
||||
|
||||
@@ -722,10 +724,12 @@ class ConfluenceConnector(
|
||||
)
|
||||
labels: list[str] = []
|
||||
if "metadata" in attachment and "labels" in attachment["metadata"]:
|
||||
for label in attachment["metadata"]["labels"].get(
|
||||
"results", []
|
||||
):
|
||||
labels.append(label.get("name", ""))
|
||||
labels.extend(
|
||||
label.get("name", "")
|
||||
for label in attachment["metadata"]["labels"].get(
|
||||
"results", []
|
||||
)
|
||||
)
|
||||
if labels:
|
||||
attachment_metadata["labels"] = labels
|
||||
page_url = page_url or build_confluence_document_id(
|
||||
@@ -1187,8 +1191,7 @@ class ConfluenceConnector(
|
||||
)
|
||||
|
||||
# Yield space hierarchy nodes first
|
||||
for node in self._yield_space_hierarchy_nodes():
|
||||
doc_metadata_list.append(node)
|
||||
doc_metadata_list.extend(self._yield_space_hierarchy_nodes())
|
||||
|
||||
# Per-page mode only: collapse shared ancestors to one GET each.
|
||||
ancestor_restrictions_cache: dict[str, dict[str, Any] | None] = {}
|
||||
@@ -1223,8 +1226,7 @@ class ConfluenceConnector(
|
||||
limit=_SLIM_DOC_BATCH_SIZE,
|
||||
):
|
||||
# Yield ancestor hierarchy nodes for this page
|
||||
for node in self._yield_ancestor_hierarchy_nodes(page):
|
||||
doc_metadata_list.append(node)
|
||||
doc_metadata_list.extend(self._yield_ancestor_hierarchy_nodes(page))
|
||||
|
||||
page_restrictions = page.get("restrictions") or {}
|
||||
page_space_key = page.get("space", {}).get("key")
|
||||
|
||||
@@ -96,20 +96,20 @@ class Document360Connector(LoadConnector, PollConnector):
|
||||
|
||||
for category in all_categories:
|
||||
if not self.categories or category["name"] in self.categories:
|
||||
for article in category["articles"]:
|
||||
articles_with_category.append(
|
||||
{"id": article["id"], "category_name": category["name"]}
|
||||
)
|
||||
articles_with_category.extend(
|
||||
{"id": article["id"], "category_name": category["name"]}
|
||||
for article in category["articles"]
|
||||
)
|
||||
for child_category in category["child_categories"]:
|
||||
all_nested_categories = flatten_child_categories(child_category)
|
||||
for nested_category in all_nested_categories:
|
||||
for article in nested_category["articles"]:
|
||||
articles_with_category.append(
|
||||
{
|
||||
"id": article["id"],
|
||||
"category_name": nested_category["name"],
|
||||
}
|
||||
)
|
||||
articles_with_category.extend(
|
||||
{
|
||||
"id": article["id"],
|
||||
"category_name": nested_category["name"],
|
||||
}
|
||||
for article in nested_category["articles"]
|
||||
)
|
||||
|
||||
return articles_with_category
|
||||
|
||||
|
||||
@@ -108,10 +108,11 @@ def _create_doc_from_transcript(transcript: dict) -> Document | None:
|
||||
meeting_organizer_email = transcript["organizer_email"]
|
||||
organizer_email_user_info = [BasicExpertInfo(email=meeting_organizer_email)]
|
||||
|
||||
meeting_participants_email_list = []
|
||||
for participant in transcript.get("participants", []):
|
||||
if participant != meeting_organizer_email and participant:
|
||||
meeting_participants_email_list.append(BasicExpertInfo(email=participant))
|
||||
meeting_participants_email_list = [
|
||||
BasicExpertInfo(email=participant)
|
||||
for participant in transcript.get("participants", [])
|
||||
if participant != meeting_organizer_email and participant
|
||||
]
|
||||
|
||||
return Document(
|
||||
id=fireflies_id,
|
||||
|
||||
@@ -427,15 +427,16 @@ class GmailConnector(
|
||||
|
||||
try:
|
||||
admin_service = get_admin_service(self.creds, self.primary_admin_email)
|
||||
emails = []
|
||||
for user in execute_paginated_retrieval(
|
||||
retrieval_function=admin_service.users().list, # ty: ignore[unresolved-attribute]
|
||||
list_key="users",
|
||||
fields=USER_FIELDS,
|
||||
domain=self.google_domain,
|
||||
):
|
||||
if email := user.get("primaryEmail"):
|
||||
emails.append(email)
|
||||
emails = [
|
||||
email
|
||||
for user in execute_paginated_retrieval(
|
||||
retrieval_function=admin_service.users().list, # ty: ignore[unresolved-attribute]
|
||||
list_key="users",
|
||||
fields=USER_FIELDS,
|
||||
domain=self.google_domain,
|
||||
)
|
||||
if (email := user.get("primaryEmail"))
|
||||
]
|
||||
return emails
|
||||
|
||||
except HttpError as e:
|
||||
|
||||
@@ -2096,14 +2096,18 @@ class GoogleDriveConnector(
|
||||
)
|
||||
|
||||
# Build slim documents
|
||||
for file in files_batch:
|
||||
if doc := build_slim_document(
|
||||
self.creds,
|
||||
file.drive_file,
|
||||
permission_sync_context,
|
||||
retriever_email=file.user_email,
|
||||
):
|
||||
slim_batch.append(doc)
|
||||
slim_batch.extend(
|
||||
doc
|
||||
for file in files_batch
|
||||
if (
|
||||
doc := build_slim_document(
|
||||
self.creds,
|
||||
file.drive_file,
|
||||
permission_sync_context,
|
||||
retriever_email=file.user_email,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Combine: hierarchy nodes first, then slim docs
|
||||
result: list[SlimDocument | HierarchyNode] = []
|
||||
|
||||
@@ -718,8 +718,9 @@ class HubSpotConnector(LoadConnector, PollConnector):
|
||||
associated_notes = self._get_associated_notes(
|
||||
api_client, ticket.id, "tickets"
|
||||
)
|
||||
for note in associated_notes:
|
||||
sections.append(self._create_object_section(note, "notes"))
|
||||
sections.extend(
|
||||
self._create_object_section(note, "notes") for note in associated_notes
|
||||
)
|
||||
|
||||
# Add association IDs to metadata
|
||||
if associated_contact_ids:
|
||||
@@ -874,8 +875,9 @@ class HubSpotConnector(LoadConnector, PollConnector):
|
||||
associated_notes = self._get_associated_notes(
|
||||
api_client, company.id, "companies"
|
||||
)
|
||||
for note in associated_notes:
|
||||
sections.append(self._create_object_section(note, "notes"))
|
||||
sections.extend(
|
||||
self._create_object_section(note, "notes") for note in associated_notes
|
||||
)
|
||||
|
||||
# Add association IDs to metadata
|
||||
if associated_contact_ids:
|
||||
@@ -1028,8 +1030,9 @@ class HubSpotConnector(LoadConnector, PollConnector):
|
||||
|
||||
# Get associated notes
|
||||
associated_notes = self._get_associated_notes(api_client, deal.id, "deals")
|
||||
for note in associated_notes:
|
||||
sections.append(self._create_object_section(note, "notes"))
|
||||
sections.extend(
|
||||
self._create_object_section(note, "notes") for note in associated_notes
|
||||
)
|
||||
|
||||
# Add association IDs to metadata
|
||||
if associated_contact_ids:
|
||||
@@ -1202,8 +1205,9 @@ class HubSpotConnector(LoadConnector, PollConnector):
|
||||
associated_notes = self._get_associated_notes(
|
||||
api_client, contact.id, "contacts"
|
||||
)
|
||||
for note in associated_notes:
|
||||
sections.append(self._create_object_section(note, "notes"))
|
||||
sections.extend(
|
||||
self._create_object_section(note, "notes") for note in associated_notes
|
||||
)
|
||||
|
||||
# Add association IDs to metadata
|
||||
if associated_company_ids:
|
||||
|
||||
@@ -955,25 +955,28 @@ class JiraConnector(
|
||||
|
||||
# Yield hierarchy nodes BEFORE the slim document (parent-before-child)
|
||||
# 1. Yield project hierarchy node (if not already yielded)
|
||||
for node in self._yield_project_hierarchy_node(
|
||||
project_key, project_name, seen_hierarchy_node_ids
|
||||
):
|
||||
slim_doc_batch.append(node)
|
||||
slim_doc_batch.extend(
|
||||
self._yield_project_hierarchy_node(
|
||||
project_key, project_name, seen_hierarchy_node_ids
|
||||
)
|
||||
)
|
||||
|
||||
# 2. If parent is an Epic, yield hierarchy node for it
|
||||
parent = best_effort_get_field_from_issue(issue, _FIELD_PARENT)
|
||||
if parent:
|
||||
for node in self._yield_parent_hierarchy_node_if_epic(
|
||||
parent, project_key, seen_hierarchy_node_ids
|
||||
):
|
||||
slim_doc_batch.append(node)
|
||||
slim_doc_batch.extend(
|
||||
self._yield_parent_hierarchy_node_if_epic(
|
||||
parent, project_key, seen_hierarchy_node_ids
|
||||
)
|
||||
)
|
||||
|
||||
# 3. If this issue IS an Epic, yield it as hierarchy node
|
||||
if self._is_epic(issue):
|
||||
for node in self._yield_epic_hierarchy_node(
|
||||
issue, project_key, seen_hierarchy_node_ids
|
||||
):
|
||||
slim_doc_batch.append(node)
|
||||
slim_doc_batch.extend(
|
||||
self._yield_epic_hierarchy_node(
|
||||
issue, project_key, seen_hierarchy_node_ids
|
||||
)
|
||||
)
|
||||
|
||||
# Now add the slim document
|
||||
issue_key = best_effort_get_field_from_issue(issue, _FIELD_KEY)
|
||||
|
||||
@@ -329,13 +329,13 @@ class LinearConnector(LoadConnector, PollConnector, OAuthConnector):
|
||||
]
|
||||
|
||||
# Add comment sections
|
||||
for comment in node["comments"]["nodes"]:
|
||||
sections.append(
|
||||
TextSection(
|
||||
link=node["url"],
|
||||
text=comment["body"] or "",
|
||||
)
|
||||
sections.extend(
|
||||
TextSection(
|
||||
link=node["url"],
|
||||
text=comment["body"] or "",
|
||||
)
|
||||
for comment in node["comments"]["nodes"]
|
||||
)
|
||||
|
||||
# Cast the sections list to the expected type
|
||||
typed_sections = cast(list[TextSection | ImageSection], sections)
|
||||
|
||||
@@ -879,7 +879,7 @@ class NotionConnector(LoadConnector, PollConnector, SlimConnector):
|
||||
page_title = None
|
||||
if hasattr(page, "database_name") and page.database_name:
|
||||
return page.database_name
|
||||
for _, prop in page.properties.items():
|
||||
for prop in page.properties.values():
|
||||
if prop["type"] == "title" and len(prop["title"]) > 0:
|
||||
page_title = " ".join([t["plain_text"] for t in prop["title"]]).strip()
|
||||
break
|
||||
|
||||
@@ -192,10 +192,9 @@ class OutlineConnector(LoadConnector, PollConnector):
|
||||
)
|
||||
|
||||
# Apply time filtering if specified
|
||||
filtered_batch: list[Document | HierarchyNode] = []
|
||||
for doc in doc_batch:
|
||||
if time_filter is None or time_filter(doc):
|
||||
filtered_batch.append(doc)
|
||||
filtered_batch: list[Document | HierarchyNode] = [
|
||||
doc for doc in doc_batch if time_filter is None or time_filter(doc)
|
||||
]
|
||||
|
||||
start_ind += num_results
|
||||
if filtered_batch:
|
||||
|
||||
@@ -104,8 +104,7 @@ def _json_to_natural_language(data: dict | list, indent: int = 0) -> str:
|
||||
else:
|
||||
result.append(f"{indent_str}{key}: {value}")
|
||||
elif isinstance(data, list):
|
||||
for item in data:
|
||||
result.append(_json_to_natural_language(item, indent + 2))
|
||||
result.extend(_json_to_natural_language(item, indent + 2) for item in data)
|
||||
|
||||
return "\n".join(result)
|
||||
|
||||
|
||||
@@ -466,7 +466,7 @@ class OnyxSalesforceSQLite:
|
||||
parent_relationship_fields = parent_relationship_fields_by_type[
|
||||
changed_type
|
||||
]
|
||||
for field_name, _ in parent_relationship_fields.items():
|
||||
for field_name in parent_relationship_fields:
|
||||
if field_name not in sf_object.data:
|
||||
logger.warning(
|
||||
"field_name=%r not in data for changed_type=%r!",
|
||||
|
||||
@@ -40,16 +40,13 @@ def get_title(soup: BeautifulSoup) -> str:
|
||||
|
||||
def get_pages(soup: BeautifulSoup, url: str) -> list[str]:
|
||||
page_tags = soup.select("li.pageNav-page")
|
||||
page_numbers = []
|
||||
for button in page_tags:
|
||||
if re.match(r"^\d+$", button.text):
|
||||
page_numbers.append(button.text)
|
||||
page_numbers = [
|
||||
button.text for button in page_tags if re.match(r"^\d+$", button.text)
|
||||
]
|
||||
|
||||
max_pages = int(max(page_numbers, key=int)) if page_numbers else 1
|
||||
|
||||
all_pages = []
|
||||
for x in range(1, int(max_pages) + 1):
|
||||
all_pages.append(f"{url}page-{x}")
|
||||
all_pages = [f"{url}page-{x}" for x in range(1, int(max_pages) + 1)]
|
||||
return all_pages
|
||||
|
||||
|
||||
|
||||
@@ -1152,23 +1152,23 @@ def slack_retrieval(
|
||||
), # ALWAYS apply exclude_channels
|
||||
}
|
||||
|
||||
for query_string in query_strings:
|
||||
search_tasks.append(
|
||||
search_tasks.extend(
|
||||
(
|
||||
query_slack,
|
||||
(
|
||||
query_slack,
|
||||
(
|
||||
query_string,
|
||||
access_token,
|
||||
query_limit,
|
||||
allowed_private_channel,
|
||||
bot_token,
|
||||
include_dm,
|
||||
dm_entities,
|
||||
available_channels,
|
||||
channel_metadata_dict,
|
||||
),
|
||||
)
|
||||
query_string,
|
||||
access_token,
|
||||
query_limit,
|
||||
allowed_private_channel,
|
||||
bot_token,
|
||||
include_dm,
|
||||
dm_entities,
|
||||
available_channels,
|
||||
channel_metadata_dict,
|
||||
),
|
||||
)
|
||||
for query_string in query_strings
|
||||
)
|
||||
|
||||
# Execute searches in parallel
|
||||
results = run_functions_tuples_in_parallel(search_tasks)
|
||||
|
||||
@@ -174,7 +174,7 @@ def merge_individual_chunks(
|
||||
chunk_to_section: dict[tuple[str, int], InferenceSection] = {}
|
||||
|
||||
# Process each document's chunks
|
||||
for _doc_id, doc_chunk_list in doc_chunks.items():
|
||||
for doc_chunk_list in doc_chunks.values():
|
||||
if not doc_chunk_list:
|
||||
continue
|
||||
|
||||
|
||||
@@ -94,8 +94,6 @@ def search_chunks(
|
||||
embedding_model: EmbeddingModel | None = None,
|
||||
prefetched_federated_retrieval_infos: list[FederatedRetrievalInfo] | None = None,
|
||||
) -> list[InferenceChunk]:
|
||||
run_queries: list[tuple[Callable, tuple]] = []
|
||||
|
||||
source_filters = (
|
||||
set(query_request.filters.source_type)
|
||||
if query_request.filters.source_type
|
||||
@@ -121,10 +119,10 @@ def search_chunks(
|
||||
federated_retrieval_info.source.to_non_federated_source()
|
||||
for federated_retrieval_info in federated_retrieval_infos
|
||||
)
|
||||
for federated_retrieval_info in federated_retrieval_infos:
|
||||
run_queries.append(
|
||||
(federated_retrieval_info.retrieval_function, (query_request,))
|
||||
)
|
||||
run_queries: list[tuple[Callable, tuple]] = [
|
||||
(federated_retrieval_info.retrieval_function, (query_request,))
|
||||
for federated_retrieval_info in federated_retrieval_infos
|
||||
]
|
||||
|
||||
# Don't run normal hybrid search if there are no indexed sources to
|
||||
# search over
|
||||
|
||||
@@ -63,14 +63,13 @@ def _relate_credential_to_user_groups__no_commit(
|
||||
credential_id: int,
|
||||
user_group_ids: list[int],
|
||||
) -> None:
|
||||
credential_user_groups = []
|
||||
for group_id in user_group_ids:
|
||||
credential_user_groups.append(
|
||||
Credential__UserGroup(
|
||||
credential_id=credential_id,
|
||||
user_group_id=group_id,
|
||||
)
|
||||
credential_user_groups = [
|
||||
Credential__UserGroup(
|
||||
credential_id=credential_id,
|
||||
user_group_id=group_id,
|
||||
)
|
||||
for group_id in user_group_ids
|
||||
]
|
||||
db_session.add_all(credential_user_groups)
|
||||
|
||||
|
||||
|
||||
@@ -398,22 +398,21 @@ class OpenSearchClient(AbstractContextManager):
|
||||
A list of IndexInfo objects for each index.
|
||||
"""
|
||||
response = self._client.cat.indices(format="json")
|
||||
indices: list[IndexInfo] = []
|
||||
for raw_index_info in response:
|
||||
indices.append(
|
||||
IndexInfo(
|
||||
name=raw_index_info.get("index", ""),
|
||||
health=raw_index_info.get("health", ""),
|
||||
status=raw_index_info.get("status", ""),
|
||||
num_primary_shards=raw_index_info.get("pri", ""),
|
||||
num_replica_shards=raw_index_info.get("rep", ""),
|
||||
docs_count=raw_index_info.get("docs.count", ""),
|
||||
docs_deleted=raw_index_info.get("docs.deleted", ""),
|
||||
created_at=raw_index_info.get("creation.date.string", ""),
|
||||
total_size=raw_index_info.get("store.size", ""),
|
||||
primary_shards_size=raw_index_info.get("pri.store.size", ""),
|
||||
)
|
||||
indices: list[IndexInfo] = [
|
||||
IndexInfo(
|
||||
name=raw_index_info.get("index", ""),
|
||||
health=raw_index_info.get("health", ""),
|
||||
status=raw_index_info.get("status", ""),
|
||||
num_primary_shards=raw_index_info.get("pri", ""),
|
||||
num_replica_shards=raw_index_info.get("rep", ""),
|
||||
docs_count=raw_index_info.get("docs.count", ""),
|
||||
docs_deleted=raw_index_info.get("docs.deleted", ""),
|
||||
created_at=raw_index_info.get("creation.date.string", ""),
|
||||
total_size=raw_index_info.get("store.size", ""),
|
||||
primary_shards_size=raw_index_info.get("pri.store.size", ""),
|
||||
)
|
||||
for raw_index_info in response
|
||||
]
|
||||
return indices
|
||||
|
||||
@log_function_time(print_only=True, debug_only=True, include_args=True)
|
||||
@@ -1377,16 +1376,15 @@ class OpenSearchIndexClient(OpenSearchClient):
|
||||
len(document_chunk_ids),
|
||||
self._index_name,
|
||||
)
|
||||
data = []
|
||||
for document_chunk_id in document_chunk_ids:
|
||||
data.append(
|
||||
{
|
||||
"_index": self._index_name,
|
||||
"_id": document_chunk_id,
|
||||
"_op_type": "update",
|
||||
"doc": properties_to_update,
|
||||
}
|
||||
)
|
||||
data = [
|
||||
{
|
||||
"_index": self._index_name,
|
||||
"_id": document_chunk_id,
|
||||
"_op_type": "update",
|
||||
"doc": properties_to_update,
|
||||
}
|
||||
for document_chunk_id in document_chunk_ids
|
||||
]
|
||||
# max_retries is the number of times to retry a request if we get a 429.
|
||||
# We do not raise on error (the default behavior of ``bulk`` is to
|
||||
# raise) because we want to attempt to retry certain failed chunks in
|
||||
|
||||
@@ -85,9 +85,10 @@ def build_vespa_filters(
|
||||
return f'"{entity}"'
|
||||
|
||||
if kg_entities:
|
||||
filter_parts = []
|
||||
for kg_entity in kg_entities:
|
||||
filter_parts.append(f"(kg_entities contains {_build_kge(kg_entity)})")
|
||||
filter_parts = [
|
||||
f"(kg_entities contains {_build_kge(kg_entity)})"
|
||||
for kg_entity in kg_entities
|
||||
]
|
||||
combined_filter_parts.append(f"({' or '.join(filter_parts)})")
|
||||
|
||||
# TODO: handle complex nested relationship logic (e.g., A participated, and B or C participated)
|
||||
|
||||
+11
-12
@@ -292,19 +292,18 @@ def _get_multi_turn_answer_with_tools(
|
||||
raise ValueError("Multi-turn eval requires 'messages' array in input")
|
||||
|
||||
# Parse messages into EvalMessage objects
|
||||
messages: list[EvalMessage] = []
|
||||
for msg_data in messages_data:
|
||||
messages.append(
|
||||
EvalMessage(
|
||||
message=msg_data["message"],
|
||||
expected_tools=msg_data.get("expected_tools", []),
|
||||
require_all_tools=msg_data.get("require_all_tools", False),
|
||||
model=msg_data.get("model"),
|
||||
model_provider=msg_data.get("model_provider"),
|
||||
temperature=msg_data.get("temperature"),
|
||||
force_tools=msg_data.get("force_tools", []),
|
||||
)
|
||||
messages: list[EvalMessage] = [
|
||||
EvalMessage(
|
||||
message=msg_data["message"],
|
||||
expected_tools=msg_data.get("expected_tools", []),
|
||||
require_all_tools=msg_data.get("require_all_tools", False),
|
||||
model=msg_data.get("model"),
|
||||
model_provider=msg_data.get("model_provider"),
|
||||
temperature=msg_data.get("temperature"),
|
||||
force_tools=msg_data.get("force_tools", []),
|
||||
)
|
||||
for msg_data in messages_data
|
||||
]
|
||||
|
||||
turn_results: list[EvalToolResult] = []
|
||||
|
||||
|
||||
@@ -67,8 +67,10 @@ def _operation_invocations(query: str) -> list[tuple[str, str]]:
|
||||
if not isinstance(definition, OperationDefinitionNode):
|
||||
continue
|
||||
operation_type = definition.operation.value # "query" | "mutation" | ...
|
||||
for field in _root_field_names(definition.selection_set, fragments, set()):
|
||||
invocations.append((operation_type, field))
|
||||
invocations.extend(
|
||||
(operation_type, field)
|
||||
for field in _root_field_names(definition.selection_set, fragments, set())
|
||||
)
|
||||
return invocations
|
||||
|
||||
|
||||
|
||||
@@ -113,13 +113,13 @@ def _split_row_by_pairs(
|
||||
current_tokens = 0
|
||||
|
||||
if pair_tokens > max_tokens:
|
||||
for split_text in split_text_by_tokens(pair_str, tokenizer, max_tokens):
|
||||
pieces.append(
|
||||
_TokenizedText(
|
||||
text=split_text,
|
||||
token_count=count_tokens(split_text, tokenizer),
|
||||
)
|
||||
pieces.extend(
|
||||
_TokenizedText(
|
||||
text=split_text,
|
||||
token_count=count_tokens(split_text, tokenizer),
|
||||
)
|
||||
for split_text in split_text_by_tokens(pair_str, tokenizer, max_tokens)
|
||||
)
|
||||
else:
|
||||
current_parts = [pair_str]
|
||||
current_tokens = pair_tokens
|
||||
|
||||
@@ -24,9 +24,10 @@ def build_total_descriptor_chunks(
|
||||
if analysis.row_count == 0:
|
||||
return []
|
||||
|
||||
lines: list[str] = []
|
||||
for idx in analysis.numeric_cols:
|
||||
lines.append(_numeric_totals_line(headers[idx], analysis.numeric_stats[idx]))
|
||||
lines: list[str] = [
|
||||
_numeric_totals_line(headers[idx], analysis.numeric_stats[idx])
|
||||
for idx in analysis.numeric_cols
|
||||
]
|
||||
for idx in analysis.categorical_cols:
|
||||
line = _categorical_top_line(headers[idx], analysis.categorical_counts[idx])
|
||||
if line:
|
||||
|
||||
@@ -105,16 +105,15 @@ def _parse_delta_tool_calls(
|
||||
if not tool_calls:
|
||||
return []
|
||||
|
||||
parsed_tool_calls: list[ChatCompletionDeltaToolCall] = []
|
||||
for tool_call in tool_calls:
|
||||
parsed_tool_calls.append(
|
||||
ChatCompletionDeltaToolCall(
|
||||
id=tool_call.get("id"),
|
||||
index=tool_call.get("index", 0),
|
||||
type=tool_call.get("type", "function"),
|
||||
function=_parse_function_call(tool_call.get("function")),
|
||||
)
|
||||
parsed_tool_calls: list[ChatCompletionDeltaToolCall] = [
|
||||
ChatCompletionDeltaToolCall(
|
||||
id=tool_call.get("id"),
|
||||
index=tool_call.get("index", 0),
|
||||
type=tool_call.get("type", "function"),
|
||||
function=_parse_function_call(tool_call.get("function")),
|
||||
)
|
||||
for tool_call in tool_calls
|
||||
]
|
||||
return parsed_tool_calls
|
||||
|
||||
|
||||
|
||||
@@ -327,12 +327,15 @@ async def _build_thread_context(
|
||||
|
||||
try:
|
||||
thread = message.channel
|
||||
messages: list[discord.Message] = []
|
||||
|
||||
# Fetch recent messages (excluding current)
|
||||
async for msg in thread.history(limit=MAX_CONTEXT_MESSAGES, oldest_first=False):
|
||||
if msg.id != message.id:
|
||||
messages.append(msg)
|
||||
messages: list[discord.Message] = [
|
||||
msg
|
||||
async for msg in thread.history(
|
||||
limit=MAX_CONTEXT_MESSAGES, oldest_first=False
|
||||
)
|
||||
if msg.id != message.id
|
||||
]
|
||||
|
||||
# Include thread starter message and its reply chain if not already present
|
||||
if thread.parent and not isinstance(thread.parent, discord.ForumChannel):
|
||||
|
||||
@@ -401,7 +401,7 @@ def get_view_values(state_values: dict[str, Any]) -> dict[str, Any]:
|
||||
dict: keys/values of the view state content
|
||||
"""
|
||||
view_values = {}
|
||||
for _, view_data in state_values.items():
|
||||
for view_data in state_values.values():
|
||||
for k, v in view_data.items():
|
||||
if (
|
||||
"selected_option" in v
|
||||
|
||||
@@ -1183,14 +1183,14 @@ def get_connector_indexing_status(
|
||||
# route is global, so a scoped manager would just 403 on click.
|
||||
federated_statuses: list[FederatedConnectorStatus] = []
|
||||
if is_connectors_admin:
|
||||
for federated_connector in federated_connectors:
|
||||
federated_statuses.append(
|
||||
FederatedConnectorStatus(
|
||||
id=federated_connector.id,
|
||||
source=federated_connector.source,
|
||||
name=f"{federated_connector.source.replace('_', ' ').title()}",
|
||||
)
|
||||
federated_statuses.extend(
|
||||
FederatedConnectorStatus(
|
||||
id=federated_connector.id,
|
||||
source=federated_connector.source,
|
||||
name=f"{federated_connector.source.replace('_', ' ').title()}",
|
||||
)
|
||||
for federated_connector in federated_connectors
|
||||
)
|
||||
|
||||
source_to_summary: dict[DocumentSource, SourceSummary] = {}
|
||||
|
||||
|
||||
@@ -603,14 +603,14 @@ class KubernetesSandboxManager(SandboxManager):
|
||||
]
|
||||
|
||||
# Add ports for session Next.js servers (one port per potential session)
|
||||
for port in range(SANDBOX_NEXTJS_PORT_START, SANDBOX_NEXTJS_PORT_END):
|
||||
ports.append(
|
||||
client.V1ServicePort(
|
||||
name=f"nextjs-{port}",
|
||||
port=port,
|
||||
target_port=port,
|
||||
)
|
||||
ports.extend(
|
||||
client.V1ServicePort(
|
||||
name=f"nextjs-{port}",
|
||||
port=port,
|
||||
target_port=port,
|
||||
)
|
||||
for port in range(SANDBOX_NEXTJS_PORT_START, SANDBOX_NEXTJS_PORT_END)
|
||||
)
|
||||
|
||||
return client.V1Service(
|
||||
api_version="v1",
|
||||
|
||||
@@ -590,8 +590,9 @@ def _render_table(
|
||||
if section_type == "table_head":
|
||||
header_cells = section.get("children", [])
|
||||
elif section_type == "table_body":
|
||||
for row in section.get("children", []):
|
||||
body_rows.append(row.get("children", []))
|
||||
body_rows.extend(
|
||||
row.get("children", []) for row in section.get("children", [])
|
||||
)
|
||||
|
||||
num_cols = len(header_cells) or (len(body_rows[0]) if body_rows else 0)
|
||||
if num_cols == 0:
|
||||
|
||||
@@ -235,10 +235,11 @@ def _extract_text_from_content(content: Any) -> str:
|
||||
if hasattr(content, "type") and content.type == "text":
|
||||
return getattr(content, "text", "") or ""
|
||||
if isinstance(content, list):
|
||||
texts = []
|
||||
for block in content:
|
||||
if hasattr(block, "type") and block.type == "text":
|
||||
texts.append(getattr(block, "text", "") or "")
|
||||
texts = [
|
||||
getattr(block, "text", "") or ""
|
||||
for block in content
|
||||
if hasattr(block, "type") and block.type == "text"
|
||||
]
|
||||
return "".join(texts)
|
||||
return ""
|
||||
|
||||
|
||||
@@ -176,16 +176,16 @@ def _run_web_search(
|
||||
list(search_results)
|
||||
)
|
||||
trimmed_results = list(filtered_results)[: request.max_results]
|
||||
for search_result in trimmed_results:
|
||||
results.append(
|
||||
LlmWebSearchResult(
|
||||
document_citation_number=DOCUMENT_CITATION_NUMBER_EMPTY_VALUE,
|
||||
url=search_result.link,
|
||||
title=search_result.title,
|
||||
snippet=search_result.snippet or "",
|
||||
unique_identifier_to_strip_away=search_result.link,
|
||||
)
|
||||
results.extend(
|
||||
LlmWebSearchResult(
|
||||
document_citation_number=DOCUMENT_CITATION_NUMBER_EMPTY_VALUE,
|
||||
url=search_result.link,
|
||||
title=search_result.title,
|
||||
snippet=search_result.snippet or "",
|
||||
unique_identifier_to_strip_away=search_result.link,
|
||||
)
|
||||
for search_result in trimmed_results
|
||||
)
|
||||
return provider_view.provider_type, results
|
||||
|
||||
|
||||
@@ -211,15 +211,14 @@ def _open_urls(
|
||||
"Web content provider failed to fetch URLs.",
|
||||
) from exc
|
||||
|
||||
results: list[LlmOpenUrlResult] = []
|
||||
for doc in docs:
|
||||
results.append(
|
||||
LlmOpenUrlResult(
|
||||
document_citation_number=DOCUMENT_CITATION_NUMBER_EMPTY_VALUE,
|
||||
content=truncate_search_result_content(doc.full_content),
|
||||
unique_identifier_to_strip_away=doc.link,
|
||||
)
|
||||
results: list[LlmOpenUrlResult] = [
|
||||
LlmOpenUrlResult(
|
||||
document_citation_number=DOCUMENT_CITATION_NUMBER_EMPTY_VALUE,
|
||||
content=truncate_search_result_content(doc.full_content),
|
||||
unique_identifier_to_strip_away=doc.link,
|
||||
)
|
||||
for doc in docs
|
||||
]
|
||||
provider_type = (
|
||||
provider_view.provider_type
|
||||
if provider_view
|
||||
|
||||
@@ -541,17 +541,14 @@ def get_federated_connector_detail(
|
||||
break
|
||||
|
||||
# Get document set mappings
|
||||
document_sets = []
|
||||
for mapping in federated_connector.document_sets:
|
||||
document_sets.append(
|
||||
{
|
||||
"id": mapping.document_set_id,
|
||||
"name": (
|
||||
mapping.document_set.name if mapping.document_set else "Unknown"
|
||||
),
|
||||
"entities": mapping.entities,
|
||||
}
|
||||
)
|
||||
document_sets = [
|
||||
{
|
||||
"id": mapping.document_set_id,
|
||||
"name": (mapping.document_set.name if mapping.document_set else "Unknown"),
|
||||
"entities": mapping.entities,
|
||||
}
|
||||
for mapping in federated_connector.document_sets
|
||||
]
|
||||
|
||||
return FederatedConnectorDetail(
|
||||
id=federated_connector.id,
|
||||
|
||||
@@ -882,19 +882,19 @@ def list_llm_provider_basics(
|
||||
|
||||
all_providers = fetch_existing_llm_providers(db_session, [])
|
||||
|
||||
accessible_providers = []
|
||||
|
||||
for provider in all_providers:
|
||||
# Use centralized access control logic with persona=None since we're
|
||||
# listing providers without a specific persona context. This correctly:
|
||||
# - Includes public providers WITHOUT persona restrictions
|
||||
# - Includes providers user can access via group membership
|
||||
# - Excludes providers with persona restrictions (requires specific persona)
|
||||
# - Excludes non-public providers with no restrictions (admin-only)
|
||||
# Use centralized access control logic with persona=None since we're
|
||||
# listing providers without a specific persona context. This correctly:
|
||||
# - Includes public providers WITHOUT persona restrictions
|
||||
# - Includes providers user can access via group membership
|
||||
# - Excludes providers with persona restrictions (requires specific persona)
|
||||
# - Excludes non-public providers with no restrictions (admin-only)
|
||||
accessible_providers = [
|
||||
LLMProviderDescriptor.from_model(provider)
|
||||
for provider in all_providers
|
||||
if can_user_access_llm_provider(
|
||||
provider, user_group_ids, persona=None, can_manage_llms=can_manage_llms
|
||||
):
|
||||
accessible_providers.append(LLMProviderDescriptor.from_model(provider))
|
||||
)
|
||||
]
|
||||
|
||||
end_time = datetime.now(timezone.utc)
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
@@ -956,9 +956,11 @@ def get_valid_model_names_for_persona(
|
||||
llm_provider_model, user_group_ids, persona, can_manage_llms=can_manage_llms
|
||||
):
|
||||
# Collect all model names from this provider
|
||||
for model_config in llm_provider_model.model_configurations:
|
||||
if model_config.is_visible:
|
||||
valid_models.append(model_config.name)
|
||||
valid_models.extend(
|
||||
model_config.name
|
||||
for model_config in llm_provider_model.model_configurations
|
||||
if model_config.is_visible
|
||||
)
|
||||
|
||||
return valid_models
|
||||
|
||||
@@ -1036,16 +1038,14 @@ def list_llm_providers_for_persona(
|
||||
db_session, [LLMModelFlowType.CHAT, LLMModelFlowType.VISION]
|
||||
)
|
||||
|
||||
llm_provider_list: list[LLMProviderDescriptor] = []
|
||||
|
||||
for llm_provider_model in all_providers:
|
||||
# Check access with persona context — respects persona restrictions
|
||||
# Check access with persona context — respects persona restrictions
|
||||
llm_provider_list: list[LLMProviderDescriptor] = [
|
||||
LLMProviderDescriptor.from_model(llm_provider_model)
|
||||
for llm_provider_model in all_providers
|
||||
if can_user_access_llm_provider(
|
||||
llm_provider_model, user_group_ids, persona, can_manage_llms=can_manage_llms
|
||||
):
|
||||
llm_provider_list.append(
|
||||
LLMProviderDescriptor.from_model(llm_provider_model)
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
end_time = datetime.now(timezone.utc)
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
@@ -432,15 +432,11 @@ def get_chat_session(
|
||||
|
||||
# Every assistant message might have a set of tool calls associated with it, these need to be replayed back for the frontend
|
||||
# Each list is the set of tool calls for the given assistant message.
|
||||
replay_packet_lists: list[list[Packet]] = []
|
||||
for msg in session_messages:
|
||||
if msg.message_type == MessageType.ASSISTANT:
|
||||
replay_packet_lists.append(
|
||||
translate_assistant_message_to_packets(
|
||||
chat_message=msg, db_session=db_session
|
||||
)
|
||||
)
|
||||
# msg_packet_list.append(Packet(ind=end_step_nr, obj=OverallStop()))
|
||||
replay_packet_lists: list[list[Packet]] = [
|
||||
translate_assistant_message_to_packets(chat_message=msg, db_session=db_session)
|
||||
for msg in session_messages
|
||||
if msg.message_type == MessageType.ASSISTANT
|
||||
]
|
||||
|
||||
return ChatSessionDetailResponse(
|
||||
chat_session_id=session_id,
|
||||
|
||||
@@ -120,16 +120,14 @@ def create_message_packets(
|
||||
def create_citation_packets(
|
||||
citation_info_list: list[CitationInfo], turn_index: int
|
||||
) -> list[Packet]:
|
||||
packets: list[Packet] = []
|
||||
|
||||
# Emit each citation as a separate CitationInfo packet
|
||||
for citation_info in citation_info_list:
|
||||
packets.append(
|
||||
Packet(
|
||||
placement=Placement(turn_index=turn_index),
|
||||
obj=citation_info,
|
||||
)
|
||||
packets: list[Packet] = [
|
||||
Packet(
|
||||
placement=Placement(turn_index=turn_index),
|
||||
obj=citation_info,
|
||||
)
|
||||
for citation_info in citation_info_list
|
||||
]
|
||||
|
||||
packets.append(Packet(placement=Placement(turn_index=turn_index), obj=SectionEnd()))
|
||||
|
||||
|
||||
@@ -93,11 +93,15 @@ def _element_text(element: dict[str, Any]) -> str:
|
||||
if shape_text:
|
||||
parts.append(_text_runs(shape_text))
|
||||
for row in element.get("table", {}).get("tableRows", []):
|
||||
for cell in row.get("tableCells", []):
|
||||
if cell.get("text"):
|
||||
parts.append(_text_runs(cell["text"]))
|
||||
for child in element.get("elementGroup", {}).get("children", []):
|
||||
parts.append(_element_text(child))
|
||||
parts.extend(
|
||||
_text_runs(cell["text"])
|
||||
for cell in row.get("tableCells", [])
|
||||
if cell.get("text")
|
||||
)
|
||||
parts.extend(
|
||||
_element_text(child)
|
||||
for child in element.get("elementGroup", {}).get("children", [])
|
||||
)
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
|
||||
@@ -301,14 +301,13 @@ class BaseSchemaValidator:
|
||||
print("PASSED - No .rels files found")
|
||||
return True
|
||||
|
||||
all_files = []
|
||||
for file_path in self.unpacked_dir.rglob("*"):
|
||||
if (
|
||||
file_path.is_file()
|
||||
and file_path.name != "[Content_Types].xml"
|
||||
and not file_path.name.endswith(".rels")
|
||||
):
|
||||
all_files.append(file_path.resolve())
|
||||
all_files = [
|
||||
file_path.resolve()
|
||||
for file_path in self.unpacked_dir.rglob("*")
|
||||
if file_path.is_file()
|
||||
and file_path.name != "[Content_Types].xml"
|
||||
and not file_path.name.endswith(".rels")
|
||||
]
|
||||
|
||||
all_referenced_files = set()
|
||||
|
||||
@@ -665,10 +664,10 @@ class BaseSchemaValidator:
|
||||
continue
|
||||
|
||||
new_errors.append(f" {relative_path}: {len(new_file_errors)} new error(s)")
|
||||
for error in list(new_file_errors)[:3]:
|
||||
new_errors.append(
|
||||
f" - {error[:250]}..." if len(error) > 250 else f" - {error}"
|
||||
)
|
||||
new_errors.extend(
|
||||
f" - {error[:250]}..." if len(error) > 250 else f" - {error}"
|
||||
for error in list(new_file_errors)[:3]
|
||||
)
|
||||
|
||||
if self.verbose:
|
||||
print(f"Validated {len(self.xml_files)} files:")
|
||||
@@ -831,10 +830,10 @@ class BaseSchemaValidator:
|
||||
return text
|
||||
matches = list(template_pattern.finditer(text))
|
||||
if matches:
|
||||
for match in matches:
|
||||
warnings.append(
|
||||
f"Found template tag in {content_type}: {match.group()}"
|
||||
)
|
||||
warnings.extend(
|
||||
f"Found template tag in {content_type}: {match.group()}"
|
||||
for match in matches
|
||||
)
|
||||
return template_pattern.sub("", text)
|
||||
return text
|
||||
|
||||
|
||||
@@ -332,20 +332,20 @@ class DOCXSchemaValidator(BaseSchemaValidator):
|
||||
}
|
||||
|
||||
orphaned_ends = range_ends - range_starts
|
||||
for comment_id in sorted(
|
||||
orphaned_ends, key=lambda x: int(x) if x and x.isdigit() else 0
|
||||
):
|
||||
errors.append(
|
||||
f' document.xml: commentRangeEnd id="{comment_id}" has no matching commentRangeStart'
|
||||
errors.extend(
|
||||
f' document.xml: commentRangeEnd id="{comment_id}" has no matching commentRangeStart'
|
||||
for comment_id in sorted(
|
||||
orphaned_ends, key=lambda x: int(x) if x and x.isdigit() else 0
|
||||
)
|
||||
)
|
||||
|
||||
orphaned_starts = range_starts - range_ends
|
||||
for comment_id in sorted(
|
||||
orphaned_starts, key=lambda x: int(x) if x and x.isdigit() else 0
|
||||
):
|
||||
errors.append(
|
||||
f' document.xml: commentRangeStart id="{comment_id}" has no matching commentRangeEnd'
|
||||
errors.extend(
|
||||
f' document.xml: commentRangeStart id="{comment_id}" has no matching commentRangeEnd'
|
||||
for comment_id in sorted(
|
||||
orphaned_starts, key=lambda x: int(x) if x and x.isdigit() else 0
|
||||
)
|
||||
)
|
||||
|
||||
comment_ids = set()
|
||||
if comments_xml and comments_xml.exists():
|
||||
@@ -359,13 +359,13 @@ class DOCXSchemaValidator(BaseSchemaValidator):
|
||||
|
||||
marker_ids = range_starts | range_ends | references
|
||||
invalid_refs = marker_ids - comment_ids
|
||||
for comment_id in sorted(
|
||||
invalid_refs, key=lambda x: int(x) if x and x.isdigit() else 0
|
||||
):
|
||||
if comment_id:
|
||||
errors.append(
|
||||
f' document.xml: marker id="{comment_id}" references non-existent comment'
|
||||
)
|
||||
errors.extend(
|
||||
f' document.xml: marker id="{comment_id}" references non-existent comment'
|
||||
for comment_id in sorted(
|
||||
invalid_refs, key=lambda x: int(x) if x and x.isdigit() else 0
|
||||
)
|
||||
if comment_id
|
||||
)
|
||||
|
||||
except (lxml.etree.XMLSyntaxError, Exception) as e:
|
||||
errors.append(f" Error parsing XML: {e}")
|
||||
|
||||
@@ -204,10 +204,11 @@ class RedliningValidator:
|
||||
author_attr = f"{{{self.namespaces['w']}}}author"
|
||||
|
||||
for parent in root.iter():
|
||||
to_remove = []
|
||||
for child in parent:
|
||||
if child.tag == ins_tag and child.get(author_attr) == self.author:
|
||||
to_remove.append(child)
|
||||
to_remove = [
|
||||
child
|
||||
for child in parent
|
||||
if child.tag == ins_tag and child.get(author_attr) == self.author
|
||||
]
|
||||
for elem in to_remove:
|
||||
parent.remove(elem)
|
||||
|
||||
@@ -215,10 +216,11 @@ class RedliningValidator:
|
||||
t_tag = f"{{{self.namespaces['w']}}}t"
|
||||
|
||||
for parent in root.iter():
|
||||
to_process = []
|
||||
for child in parent:
|
||||
if child.tag == del_tag and child.get(author_attr) == self.author:
|
||||
to_process.append((child, list(parent).index(child)))
|
||||
to_process = [
|
||||
(child, list(parent).index(child))
|
||||
for child in parent
|
||||
if child.tag == del_tag and child.get(author_attr) == self.author
|
||||
]
|
||||
|
||||
for del_elem, del_index in reversed(to_process):
|
||||
for elem in del_elem.iter():
|
||||
@@ -235,10 +237,9 @@ class RedliningValidator:
|
||||
|
||||
paragraphs = []
|
||||
for p_elem in root.findall(f".//{p_tag}"):
|
||||
text_parts = []
|
||||
for t_elem in p_elem.findall(f".//{t_tag}"):
|
||||
if t_elem.text:
|
||||
text_parts.append(t_elem.text)
|
||||
text_parts = [
|
||||
t_elem.text for t_elem in p_elem.findall(f".//{t_tag}") if t_elem.text
|
||||
]
|
||||
paragraph_text = "".join(text_parts)
|
||||
if paragraph_text:
|
||||
paragraphs.append(paragraph_text)
|
||||
|
||||
@@ -305,7 +305,7 @@ def _resolve_urls_to_document_ids(
|
||||
def _estimate_result_chars(result: dict[str, Any]) -> int:
|
||||
"""Estimate character count from document fields in a result dict."""
|
||||
total = 0
|
||||
for _key, value in result.items():
|
||||
for value in result.values():
|
||||
if value is not None:
|
||||
total += len(str(value))
|
||||
return total
|
||||
|
||||
@@ -114,8 +114,7 @@ def _normalize_text_with_mapping(text: str) -> tuple[str, list[int]]:
|
||||
nfd_to_orig: list[int] = []
|
||||
for orig_idx, orig_char in enumerate(original_text):
|
||||
nfd_of_char = unicodedata.normalize("NFD", orig_char)
|
||||
for _ in nfd_of_char:
|
||||
nfd_to_orig.append(orig_idx)
|
||||
nfd_to_orig.extend([orig_idx] * len(nfd_of_char))
|
||||
|
||||
# Map NFC positions → NFD positions.
|
||||
# Each NFC char, when decomposed, tells us exactly how many NFD
|
||||
|
||||
@@ -10,8 +10,9 @@ def filter_web_contents_with_no_title_or_content(
|
||||
Downstream uses these fields for display + prompting; drop empty ones centrally
|
||||
rather than duplicating checks across provider clients.
|
||||
"""
|
||||
filtered: list[WebContent] = []
|
||||
for content in contents:
|
||||
if content.title.strip() or content.full_content.strip():
|
||||
filtered.append(content)
|
||||
filtered: list[WebContent] = [
|
||||
content
|
||||
for content in contents
|
||||
if content.title.strip() or content.full_content.strip()
|
||||
]
|
||||
return filtered
|
||||
@@ -923,12 +923,12 @@ class SearchTool(Tool[SearchToolOverrideKwargs]):
|
||||
if semantic_query
|
||||
else []
|
||||
)
|
||||
for llm_query in llm_queries:
|
||||
# In rare cases, the LLM may fail to provide real queries
|
||||
if llm_query:
|
||||
semantic_queries_with_weights.append(
|
||||
(llm_query, LLM_NON_CUSTOM_QUERY_WEIGHT)
|
||||
)
|
||||
# In rare cases, the LLM may fail to provide real queries
|
||||
semantic_queries_with_weights.extend(
|
||||
(llm_query, LLM_NON_CUSTOM_QUERY_WEIGHT)
|
||||
for llm_query in llm_queries
|
||||
if llm_query
|
||||
)
|
||||
if override_kwargs.original_query:
|
||||
semantic_queries_with_weights.append(
|
||||
(override_kwargs.original_query, ORIGINAL_QUERY_WEIGHT)
|
||||
|
||||
@@ -243,7 +243,7 @@ def merge_overlapping_sections(
|
||||
merged_sections: dict[tuple[str, int], InferenceSection] = {}
|
||||
|
||||
# Process each document's sections
|
||||
for _doc_id, doc_section_list in doc_sections.items():
|
||||
for doc_section_list in doc_sections.values():
|
||||
if not doc_section_list:
|
||||
continue
|
||||
|
||||
|
||||
@@ -24,10 +24,9 @@ def filter_web_search_results_with_no_title_or_snippet(
|
||||
titles/snippets for display and prompting, so we drop those empty entries
|
||||
centrally (rather than duplicating the check in each client).
|
||||
"""
|
||||
filtered: list[WebSearchResult] = []
|
||||
for result in results:
|
||||
if result.title.strip() or result.snippet.strip():
|
||||
filtered.append(result)
|
||||
filtered: list[WebSearchResult] = [
|
||||
result for result in results if result.title.strip() or result.snippet.strip()
|
||||
]
|
||||
return filtered
|
||||
|
||||
|
||||
|
||||
@@ -88,10 +88,11 @@ def _build_route_map(app: FastAPI) -> list[tuple[re.Pattern[str], str]]:
|
||||
Used by endpoint context middleware to resolve request paths to route
|
||||
templates, avoiding high-cardinality raw paths in metrics labels.
|
||||
"""
|
||||
route_map: list[tuple[re.Pattern[str], str]] = []
|
||||
for route in app.routes:
|
||||
if isinstance(route, APIRoute):
|
||||
route_map.append((route.path_regex, route.path))
|
||||
route_map: list[tuple[re.Pattern[str], str]] = [
|
||||
(route.path_regex, route.path)
|
||||
for route in app.routes
|
||||
if isinstance(route, APIRoute)
|
||||
]
|
||||
return route_map
|
||||
|
||||
|
||||
|
||||
@@ -76,9 +76,7 @@ def generate_dummy_chunk(
|
||||
image_file_id=None,
|
||||
)
|
||||
|
||||
document_set_names = []
|
||||
for i in range(number_of_document_sets):
|
||||
document_set_names.append(f"Document Set {i}")
|
||||
document_set_names = [f"Document Set {i}" for i in range(number_of_document_sets)]
|
||||
|
||||
user_emails: list[str | None] = []
|
||||
user_groups: list[str] = []
|
||||
|
||||
@@ -132,8 +132,7 @@ def collect_control_plane_data() -> list[dict[str, Any]]:
|
||||
control_plane_data = []
|
||||
with open("control_plane_data.csv", "r", newline="", encoding="utf-8") as csvfile:
|
||||
reader = csv.DictReader(csvfile)
|
||||
for row in reader:
|
||||
control_plane_data.append(row)
|
||||
control_plane_data.extend(reader)
|
||||
|
||||
return control_plane_data
|
||||
|
||||
|
||||
@@ -345,9 +345,7 @@ def test_airtable_connector_index_all(
|
||||
|
||||
all_docs: list[Document] = []
|
||||
for batch in connector.load_from_state():
|
||||
for item in batch:
|
||||
if isinstance(item, Document):
|
||||
all_docs.append(item)
|
||||
all_docs.extend(item for item in batch if isinstance(item, Document))
|
||||
|
||||
# 2 from Tickets + 4 from Support Categories + 1 from Table 3 = 7
|
||||
assert len(all_docs) == 7
|
||||
|
||||
@@ -258,9 +258,7 @@ class TestLoadFromState:
|
||||
reference_data: dict[str, Any], # noqa: ARG002
|
||||
) -> None:
|
||||
"""Test that no documents are yielded twice."""
|
||||
document_ids = []
|
||||
for doc in connector_doc_generator(connector):
|
||||
document_ids.append(doc.id)
|
||||
document_ids = [doc.id for doc in connector_doc_generator(connector)]
|
||||
|
||||
unique_ids = set(document_ids)
|
||||
assert len(document_ids) == len(unique_ids), (
|
||||
|
||||
@@ -28,10 +28,9 @@ def test_discord_connector_basic(discord_connector: DiscordConnector) -> None:
|
||||
|
||||
doc_batch = next(doc_batch_generator)
|
||||
|
||||
docs: list[Document] = []
|
||||
for doc in doc_batch:
|
||||
if not isinstance(doc, HierarchyNode):
|
||||
docs.append(doc)
|
||||
docs: list[Document] = [
|
||||
doc for doc in doc_batch if not isinstance(doc, HierarchyNode)
|
||||
]
|
||||
|
||||
assert len(docs) > 0, "No documents were retrieved from the connector"
|
||||
|
||||
|
||||
@@ -77,12 +77,12 @@ def test_answer_with_only_anthropic_provider(
|
||||
chat_session_id=chat_session.id,
|
||||
)
|
||||
|
||||
response_stream: list[AnswerStreamPart] = []
|
||||
for packet in handle_stream_message_objects(
|
||||
new_msg_req=chat_request,
|
||||
user=test_user,
|
||||
):
|
||||
response_stream.append(packet)
|
||||
response_stream: list[AnswerStreamPart] = list(
|
||||
handle_stream_message_objects(
|
||||
new_msg_req=chat_request,
|
||||
user=test_user,
|
||||
)
|
||||
)
|
||||
|
||||
assert response_stream, "Should receive streamed packets"
|
||||
assert not any(
|
||||
|
||||
@@ -193,8 +193,7 @@ class TestClaimConcurrency:
|
||||
results: list[bool] = []
|
||||
with ThreadPoolExecutor(max_workers=num_threads) as executor:
|
||||
futures = [executor.submit(claim) for _ in range(num_threads)]
|
||||
for future in as_completed(futures):
|
||||
results.append(future.result())
|
||||
results.extend(future.result() for future in as_completed(futures))
|
||||
|
||||
winners = sum(1 for r in results if r)
|
||||
assert winners == 1, f"Expected 1 winner, got {winners}"
|
||||
|
||||
+2
-3
@@ -435,9 +435,8 @@ class TestPerformExternalGroupSync:
|
||||
def test_batch_processing(self, db_session: Session) -> None:
|
||||
"""Test that large numbers of groups are processed in batches"""
|
||||
# Create many test users
|
||||
users = []
|
||||
for i in range(150): # More than the batch size of 100
|
||||
users.append(_create_ext_perm_user(db_session, f"user{i}"))
|
||||
# More than the batch size of 100
|
||||
users = [_create_ext_perm_user(db_session, f"user{i}") for i in range(150)]
|
||||
|
||||
cc_pair = _create_test_connector_credential_pair(db_session)
|
||||
|
||||
|
||||
@@ -249,11 +249,12 @@ class TestTagRaceCondition:
|
||||
|
||||
# Run both types of operations concurrently
|
||||
with ThreadPoolExecutor(max_workers=num_documents * 2) as executor:
|
||||
futures: list[Future[Union[Tag | None] | list[Tag]]] = []
|
||||
for doc_id in doc_ids_single:
|
||||
futures.append(executor.submit(create_single_tag, doc_id))
|
||||
for doc_id in doc_ids_list:
|
||||
futures.append(executor.submit(create_list_tag, doc_id))
|
||||
futures: list[Future[Union[Tag | None] | list[Tag]]] = [
|
||||
executor.submit(create_single_tag, doc_id) for doc_id in doc_ids_single
|
||||
]
|
||||
futures.extend(
|
||||
executor.submit(create_list_tag, doc_id) for doc_id in doc_ids_list
|
||||
)
|
||||
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
|
||||
@@ -48,16 +48,15 @@ def test_upsert_parents_handles_deep_chain(
|
||||
"""A 1500-deep parent chain must upsert without RecursionError."""
|
||||
tag = uuid4().hex[:8]
|
||||
# Build chain: root -> n1 -> n2 -> ... -> n{DEPTH-1}
|
||||
nodes: list[PydanticHierarchyNode] = []
|
||||
for i in range(DEEP_CHAIN_DEPTH):
|
||||
nodes.append(
|
||||
PydanticHierarchyNode(
|
||||
raw_node_id=f"deep_{tag}_{i}",
|
||||
raw_parent_id=f"deep_{tag}_{i - 1}" if i > 0 else None,
|
||||
display_name=f"Deep {i}",
|
||||
node_type=HierarchyNodeType.PAGE,
|
||||
)
|
||||
nodes: list[PydanticHierarchyNode] = [
|
||||
PydanticHierarchyNode(
|
||||
raw_node_id=f"deep_{tag}_{i}",
|
||||
raw_parent_id=f"deep_{tag}_{i - 1}" if i > 0 else None,
|
||||
display_name=f"Deep {i}",
|
||||
node_type=HierarchyNodeType.PAGE,
|
||||
)
|
||||
for i in range(DEEP_CHAIN_DEPTH)
|
||||
]
|
||||
|
||||
# Pass the deepest child first so upsert_parents has to walk the entire
|
||||
# chain in one call (this is the worst-case shape that originally blew the
|
||||
|
||||
@@ -527,13 +527,9 @@ class TestDefaultProviderEndpoint:
|
||||
existing_providers = fetch_existing_llm_providers(
|
||||
db_session, flow_type_filter=[LLMModelFlowType.CHAT]
|
||||
)
|
||||
provider_names_to_restore: list[str] = []
|
||||
|
||||
for provider in existing_providers:
|
||||
if provider.name is not None:
|
||||
provider_names_to_restore.append(provider.name)
|
||||
|
||||
# Remove all providers temporarily
|
||||
# Remove all providers temporarily. The `finally` rollback restores
|
||||
# them, since none of these deletes are committed.
|
||||
for provider in existing_providers:
|
||||
remove_llm_provider(db_session, provider.id)
|
||||
|
||||
|
||||
+3
-2
@@ -165,8 +165,9 @@ def test_thread_safe_concurrent_reads() -> None:
|
||||
|
||||
def worker() -> None:
|
||||
try:
|
||||
for _ in range(50):
|
||||
results.append(get_security_settings().user_directory_admin_only)
|
||||
results.extend(
|
||||
get_security_settings().user_directory_admin_only for _ in range(50)
|
||||
)
|
||||
except BaseException as e: # noqa: BLE001
|
||||
errors.append(e)
|
||||
|
||||
|
||||
@@ -199,8 +199,9 @@ class CCPairManager:
|
||||
indexing_statuses = []
|
||||
for connectors_by_source in indexing_status_response:
|
||||
connectors = connectors_by_source["indexing_statuses"]
|
||||
for connector in connectors:
|
||||
indexing_statuses.append(ConnectorIndexingStatusLite(**connector))
|
||||
indexing_statuses.extend(
|
||||
ConnectorIndexingStatusLite(**connector) for connector in connectors
|
||||
)
|
||||
return indexing_statuses
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -43,16 +43,15 @@ class FileManager:
|
||||
|
||||
response_json = response.json()
|
||||
# Convert UserFileSnapshot to FileDescriptor format
|
||||
file_descriptors: List[FileDescriptor] = []
|
||||
for user_file in response_json.get("user_files", []):
|
||||
file_descriptors.append(
|
||||
{
|
||||
"id": user_file["file_id"],
|
||||
"type": user_file["chat_file_type"],
|
||||
"name": user_file["name"],
|
||||
"user_file_id": str(user_file["id"]),
|
||||
}
|
||||
)
|
||||
file_descriptors: List[FileDescriptor] = [
|
||||
{
|
||||
"id": user_file["file_id"],
|
||||
"type": user_file["chat_file_type"],
|
||||
"name": user_file["name"],
|
||||
"user_file_id": str(user_file["id"]),
|
||||
}
|
||||
for user_file in response_json.get("user_files", [])
|
||||
]
|
||||
return file_descriptors, ""
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -100,8 +100,7 @@ def test_index_attempt_pagination(reset: None) -> None: # noqa: ARG001
|
||||
base_time=base_time,
|
||||
)
|
||||
|
||||
for attempt in generated_attempts:
|
||||
all_attempt_ids.append(attempt.id)
|
||||
all_attempt_ids.extend(attempt.id for attempt in generated_attempts)
|
||||
|
||||
# Verify basic pagination with different page sizes
|
||||
print("Verifying basic pagination with page size 5")
|
||||
|
||||
@@ -60,7 +60,7 @@ def test_query_history_pagination(reset: None) -> None: # noqa: ARG001
|
||||
) = setup_chat_sessions_with_different_feedback()
|
||||
|
||||
all_chat_sessions = []
|
||||
for _, chat_sessions in chat_sessions_by_feedback_type.items():
|
||||
for chat_sessions in chat_sessions_by_feedback_type.values():
|
||||
all_chat_sessions.extend(chat_sessions)
|
||||
|
||||
# Verify basic pagination with different page sizes
|
||||
|
||||
@@ -248,11 +248,12 @@ def get_api_server_host_port(env_name: str) -> str:
|
||||
|
||||
stdout, _ = _run_command("docker ps -a --format '{{json .}}'")
|
||||
containers = [json.loads(line) for line in stdout.splitlines()]
|
||||
server_jsons = []
|
||||
|
||||
for container in containers:
|
||||
if container_name in container["Names"] and env_name in container["Names"]:
|
||||
server_jsons.append(container)
|
||||
server_jsons = [
|
||||
container
|
||||
for container in containers
|
||||
if container_name in container["Names"] and env_name in container["Names"]
|
||||
]
|
||||
|
||||
if not server_jsons:
|
||||
raise RuntimeError(
|
||||
|
||||
@@ -25,8 +25,7 @@ def unzip_and_get_file_paths(zip_file_path: str) -> list[str]:
|
||||
|
||||
file_paths = []
|
||||
for root, _, files in os.walk(persistent_dir):
|
||||
for file in sorted(files):
|
||||
file_paths.append(os.path.join(root, file))
|
||||
file_paths.extend(os.path.join(root, file) for file in sorted(files))
|
||||
|
||||
return file_paths
|
||||
|
||||
|
||||
@@ -68,8 +68,7 @@ def _stream_fragments(
|
||||
for packet in maybe_emit_argument_delta(tc_map, delta, pl, parsers=parsers):
|
||||
obj = packet.obj
|
||||
assert isinstance(obj, ToolCallArgumentDelta)
|
||||
for value in obj.argument_deltas.values():
|
||||
emitted.append(value)
|
||||
emitted.extend(obj.argument_deltas.values())
|
||||
return emitted
|
||||
|
||||
|
||||
|
||||
@@ -358,8 +358,7 @@ def test_formatted_citation_yielded_separately(
|
||||
|
||||
results = []
|
||||
for token in ["Text [", "1", "] here."]:
|
||||
for result in processor.process_token(token):
|
||||
results.append(result)
|
||||
results.extend(processor.process_token(token))
|
||||
|
||||
# Should have text chunks and formatted citation
|
||||
text_results = [r for r in results if isinstance(r, str)]
|
||||
@@ -887,12 +886,10 @@ def test_stop_token_detection_stops_processing() -> None:
|
||||
|
||||
results = []
|
||||
for token in ["Text ", "ST", "OP"]:
|
||||
for result in processor.process_token(token):
|
||||
results.append(result)
|
||||
results.extend(processor.process_token(token))
|
||||
|
||||
# Try to add more text after stop token
|
||||
for result in processor.process_token(" more text"):
|
||||
results.append(result)
|
||||
results.extend(processor.process_token(" more text"))
|
||||
|
||||
# Processing should stop at STOP token - no results after STOP
|
||||
output = "".join(r for r in results if isinstance(r, str))
|
||||
@@ -907,8 +904,7 @@ def test_partial_stop_token_held_back() -> None:
|
||||
|
||||
results = []
|
||||
for token in ["Text ", "ST"]:
|
||||
for result in processor.process_token(token):
|
||||
results.append(result)
|
||||
results.extend(processor.process_token(token))
|
||||
|
||||
# Partial stop token should be held back
|
||||
output = "".join(r for r in results if isinstance(r, str))
|
||||
@@ -924,8 +920,7 @@ def test_stop_token_at_different_positions() -> None:
|
||||
processor1 = DynamicCitationProcessor(stop_stream=stop_stream)
|
||||
results1 = []
|
||||
for token in ["END"]:
|
||||
for result in processor1.process_token(token):
|
||||
results1.append(result)
|
||||
results1.extend(processor1.process_token(token))
|
||||
# Stop token detection returns early, so no results
|
||||
output1 = "".join(r for r in results1 if isinstance(r, str))
|
||||
assert output1 == "" # Stop token detected, no output
|
||||
@@ -934,8 +929,7 @@ def test_stop_token_at_different_positions() -> None:
|
||||
processor2 = DynamicCitationProcessor(stop_stream=stop_stream)
|
||||
results2 = []
|
||||
for token in ["Start ", "EN", "D"]:
|
||||
for result in processor2.process_token(token):
|
||||
results2.append(result)
|
||||
results2.extend(processor2.process_token(token))
|
||||
output2 = "".join(r for r in results2 if isinstance(r, str))
|
||||
# "Start " should be processed before stop token is detected
|
||||
assert "Start " in output2
|
||||
@@ -967,12 +961,10 @@ def test_none_token_flushes_remaining_segment(
|
||||
|
||||
results = []
|
||||
for token in ["Remaining ", "text"]:
|
||||
for result in processor.process_token(token):
|
||||
results.append(result)
|
||||
results.extend(processor.process_token(token))
|
||||
|
||||
# Flush with None
|
||||
for result in processor.process_token(None):
|
||||
results.append(result)
|
||||
results.extend(processor.process_token(None))
|
||||
|
||||
output = "".join(r for r in results if isinstance(r, str))
|
||||
assert "Remaining text" in output
|
||||
|
||||
@@ -179,9 +179,7 @@ SAMPLE_BASES = [
|
||||
def _collect_docs(connector: AirtableConnector) -> list[Document]:
|
||||
docs: list[Document] = []
|
||||
for batch in connector.load_from_state():
|
||||
for item in batch:
|
||||
if isinstance(item, Document):
|
||||
docs.append(item)
|
||||
docs.extend(item for item in batch if isinstance(item, Document))
|
||||
return docs
|
||||
|
||||
|
||||
|
||||
@@ -81,9 +81,7 @@ def _collect_slim_doc_ids(
|
||||
):
|
||||
ids: list[str] = []
|
||||
for batch in connector.retrieve_all_slim_docs():
|
||||
for item in batch:
|
||||
if isinstance(item, SlimDocument):
|
||||
ids.append(item.id)
|
||||
ids.extend(item.id for item in batch if isinstance(item, SlimDocument))
|
||||
return ids
|
||||
|
||||
|
||||
|
||||
@@ -98,9 +98,7 @@ def _collect_slim_doc_ids(connector: ConfluenceConnector) -> list[str]:
|
||||
):
|
||||
ids: list[str] = []
|
||||
for batch in connector.retrieve_all_slim_docs():
|
||||
for item in batch:
|
||||
if isinstance(item, SlimDocument):
|
||||
ids.append(item.id)
|
||||
ids.extend(item.id for item in batch if isinstance(item, SlimDocument))
|
||||
return ids
|
||||
|
||||
|
||||
|
||||
@@ -101,15 +101,15 @@ def compute_uncovered_units(
|
||||
if spec.untested is not None:
|
||||
continue
|
||||
for variant in spec.variants or (None,):
|
||||
for capability in spec.capabilities:
|
||||
if (operation, variant, capability) not in exercised:
|
||||
uncovered.append(
|
||||
UncoveredUnit(
|
||||
operation=operation,
|
||||
variant=variant,
|
||||
capability=capability,
|
||||
)
|
||||
)
|
||||
uncovered.extend(
|
||||
UncoveredUnit(
|
||||
operation=operation,
|
||||
variant=variant,
|
||||
capability=capability,
|
||||
)
|
||||
for capability in spec.capabilities
|
||||
if (operation, variant, capability) not in exercised
|
||||
)
|
||||
return uncovered
|
||||
|
||||
|
||||
@@ -171,13 +171,13 @@ def find_import_fence_violations(
|
||||
imported = [node.module]
|
||||
else:
|
||||
continue
|
||||
for module in imported:
|
||||
if module.split(".")[0] in roots:
|
||||
violations.append(
|
||||
FenceViolation(
|
||||
file=str(path),
|
||||
line=node.lineno,
|
||||
module=module,
|
||||
)
|
||||
)
|
||||
violations.extend(
|
||||
FenceViolation(
|
||||
file=str(path),
|
||||
line=node.lineno,
|
||||
module=module,
|
||||
)
|
||||
for module in imported
|
||||
if module.split(".")[0] in roots
|
||||
)
|
||||
return violations
|
||||
@@ -183,8 +183,7 @@ def test_multiple_docs_each_under_limit_flush_per_doc() -> None:
|
||||
chunks = []
|
||||
for doc_idx in range(3):
|
||||
doc_id = f"doc_{doc_idx}"
|
||||
for chunk_idx in range(50):
|
||||
chunks.append(_make_chunk(doc_id, chunk_idx))
|
||||
chunks.extend(_make_chunk(doc_id, chunk_idx) for chunk_idx in range(50))
|
||||
|
||||
metadata = IndexingMetadata(
|
||||
doc_id_to_chunk_cnt_diff={
|
||||
|
||||
@@ -262,8 +262,7 @@ class TestEmbedChunksInBatches:
|
||||
def _embed(
|
||||
chunks: list[DocAwareChunk], **_kwargs: object
|
||||
) -> tuple[list[IndexChunk], list[ConnectorFailure]]:
|
||||
for c in chunks:
|
||||
embedded_doc_ids.append(c.source_document.id)
|
||||
embedded_doc_ids.extend(c.source_document.id for c in chunks)
|
||||
return _mock_embed_fail_doc("docA")(chunks)
|
||||
|
||||
mock_embed.side_effect = _embed
|
||||
|
||||
+12
-11
@@ -128,17 +128,18 @@ def _run_send_message(
|
||||
events: list[Any] = []
|
||||
|
||||
def runner() -> None:
|
||||
for evt in client.send_message(
|
||||
_SESSION,
|
||||
"hello",
|
||||
directory=_DIRECTORY,
|
||||
model_provider=model_provider,
|
||||
model_id=model_id,
|
||||
attachments=attachments,
|
||||
timeout=timeout,
|
||||
absolute_timeout=absolute_timeout,
|
||||
):
|
||||
events.append(evt)
|
||||
events.extend(
|
||||
client.send_message(
|
||||
_SESSION,
|
||||
"hello",
|
||||
directory=_DIRECTORY,
|
||||
model_provider=model_provider,
|
||||
model_id=model_id,
|
||||
attachments=attachments,
|
||||
timeout=timeout,
|
||||
absolute_timeout=absolute_timeout,
|
||||
)
|
||||
)
|
||||
|
||||
t = threading.Thread(target=runner, daemon=True)
|
||||
t.start()
|
||||
|
||||
@@ -263,16 +263,15 @@ def test_table_cell_alignment_follows_column_spec() -> None:
|
||||
# The Markdown separator row (:--, :--:, --:) sets per-column alignment, the
|
||||
# way pandoc renders it.
|
||||
doc = _render("| L | C | R |\n|:--|:--:|--:|\n| a | b | c |\n")
|
||||
row_alignments = []
|
||||
for row in doc.tables[0].rows:
|
||||
row_alignments.append(
|
||||
[
|
||||
None
|
||||
if cell.paragraphs[0].alignment is None
|
||||
else cell.paragraphs[0].alignment
|
||||
for cell in row.cells
|
||||
]
|
||||
)
|
||||
row_alignments = [
|
||||
[
|
||||
None
|
||||
if cell.paragraphs[0].alignment is None
|
||||
else cell.paragraphs[0].alignment
|
||||
for cell in row.cells
|
||||
]
|
||||
for row in doc.tables[0].rows
|
||||
]
|
||||
# WD_ALIGN_PARAGRAPH: left default (None), CENTER == 1, RIGHT == 2.
|
||||
assert all(
|
||||
alignments[1] == WD_ALIGN_PARAGRAPH.CENTER for alignments in row_alignments
|
||||
|
||||
@@ -32,8 +32,7 @@ def _collect_with_timeout(
|
||||
|
||||
def runner() -> None:
|
||||
try:
|
||||
for item in gen:
|
||||
items.append(item)
|
||||
items.extend(gen)
|
||||
except BaseException as e: # noqa: BLE001
|
||||
error.append(e)
|
||||
|
||||
|
||||
@@ -29,10 +29,11 @@ from onyx.utils.audit import (
|
||||
|
||||
def _capture(caplog: pytest.LogCaptureFixture) -> list[dict[str, Any]]:
|
||||
"""Parse every captured ``onyx.audit`` JSON message into a dict."""
|
||||
events: list[dict[str, Any]] = []
|
||||
for record in caplog.records:
|
||||
if record.name.startswith("onyx.audit"):
|
||||
events.append(json.loads(record.getMessage()))
|
||||
events: list[dict[str, Any]] = [
|
||||
json.loads(record.getMessage())
|
||||
for record in caplog.records
|
||||
if record.name.startswith("onyx.audit")
|
||||
]
|
||||
return events
|
||||
|
||||
|
||||
|
||||
@@ -123,58 +123,65 @@ class TestEscapeSequences:
|
||||
|
||||
def test_newline_escape(self) -> None:
|
||||
deltas = _all_deltas(['{"text": "line1\\nline2"}'])
|
||||
text_parts = []
|
||||
for d in deltas:
|
||||
if isinstance(d, dict) and "text" in d and isinstance(d["text"], str):
|
||||
text_parts.append(d["text"])
|
||||
text_parts = [
|
||||
d["text"]
|
||||
for d in deltas
|
||||
if isinstance(d, dict) and "text" in d and isinstance(d["text"], str)
|
||||
]
|
||||
assert "".join(text_parts) == "line1\nline2"
|
||||
|
||||
def test_tab_escape(self) -> None:
|
||||
deltas = _all_deltas(['{"t": "a\\tb"}'])
|
||||
parts = []
|
||||
for d in deltas:
|
||||
if isinstance(d, dict) and "t" in d and isinstance(d["t"], str):
|
||||
parts.append(d["t"])
|
||||
parts = [
|
||||
d["t"]
|
||||
for d in deltas
|
||||
if isinstance(d, dict) and "t" in d and isinstance(d["t"], str)
|
||||
]
|
||||
assert "".join(parts) == "a\tb"
|
||||
|
||||
def test_escaped_quote(self) -> None:
|
||||
deltas = _all_deltas(['{"q": "say \\"hi\\""}'])
|
||||
parts = []
|
||||
for d in deltas:
|
||||
if isinstance(d, dict) and "q" in d and isinstance(d["q"], str):
|
||||
parts.append(d["q"])
|
||||
parts = [
|
||||
d["q"]
|
||||
for d in deltas
|
||||
if isinstance(d, dict) and "q" in d and isinstance(d["q"], str)
|
||||
]
|
||||
assert "".join(parts) == 'say "hi"'
|
||||
|
||||
def test_unicode_escape(self) -> None:
|
||||
deltas = _all_deltas(['{"u": "\\u0041\\u0042"}'])
|
||||
parts = []
|
||||
for d in deltas:
|
||||
if isinstance(d, dict) and "u" in d and isinstance(d["u"], str):
|
||||
parts.append(d["u"])
|
||||
parts = [
|
||||
d["u"]
|
||||
for d in deltas
|
||||
if isinstance(d, dict) and "u" in d and isinstance(d["u"], str)
|
||||
]
|
||||
assert "".join(parts) == "AB"
|
||||
|
||||
def test_escape_split_across_chunks(self) -> None:
|
||||
deltas = _all_deltas(['{"x": "a\\', 'nb"}'])
|
||||
parts = []
|
||||
for d in deltas:
|
||||
if isinstance(d, dict) and "x" in d and isinstance(d["x"], str):
|
||||
parts.append(d["x"])
|
||||
parts = [
|
||||
d["x"]
|
||||
for d in deltas
|
||||
if isinstance(d, dict) and "x" in d and isinstance(d["x"], str)
|
||||
]
|
||||
assert "".join(parts) == "a\nb"
|
||||
|
||||
def test_unicode_escape_split_across_chunks(self) -> None:
|
||||
deltas = _all_deltas(['{"u": "\\u00', '41"}'])
|
||||
parts = []
|
||||
for d in deltas:
|
||||
if isinstance(d, dict) and "u" in d and isinstance(d["u"], str):
|
||||
parts.append(d["u"])
|
||||
parts = [
|
||||
d["u"]
|
||||
for d in deltas
|
||||
if isinstance(d, dict) and "u" in d and isinstance(d["u"], str)
|
||||
]
|
||||
assert "".join(parts) == "A"
|
||||
|
||||
def test_backslash_escape(self) -> None:
|
||||
deltas = _all_deltas(['{"p": "c:\\\\dir"}'])
|
||||
parts = []
|
||||
for d in deltas:
|
||||
if isinstance(d, dict) and "p" in d and isinstance(d["p"], str):
|
||||
parts.append(d["p"])
|
||||
parts = [
|
||||
d["p"]
|
||||
for d in deltas
|
||||
if isinstance(d, dict) and "p" in d and isinstance(d["p"], str)
|
||||
]
|
||||
assert "".join(parts) == "c:\\dir"
|
||||
|
||||
|
||||
@@ -196,9 +203,7 @@ class TestNestedStructures:
|
||||
all_items: list[str] = []
|
||||
for d in deltas:
|
||||
if isinstance(d, list):
|
||||
for item in d:
|
||||
if isinstance(item, str):
|
||||
all_items.append(item)
|
||||
all_items.extend(item for item in d if isinstance(item, str))
|
||||
elif isinstance(d, str):
|
||||
all_items.append(d)
|
||||
joined = "".join(all_items)
|
||||
|
||||
@@ -304,11 +304,12 @@ def test_parallel_yield_basic() -> None:
|
||||
gen3 = make_gen([3, 6, 9], 0.15) # Slowest generator
|
||||
|
||||
# Collect results with timestamps
|
||||
results: list[tuple[float, int]] = []
|
||||
start_time = time.time()
|
||||
|
||||
for value in parallel_yield([gen1, gen2, gen3]):
|
||||
results.append((time.time() - start_time, value))
|
||||
results: list[tuple[float, int]] = [
|
||||
(time.time() - start_time, value)
|
||||
for value in parallel_yield([gen1, gen2, gen3])
|
||||
]
|
||||
|
||||
# Verify all values were yielded
|
||||
assert sorted(v for _, v in results) == list(range(1, 10))
|
||||
|
||||
@@ -46,7 +46,7 @@ def test_run_functions_in_parallel_preserves_contextvar() -> None:
|
||||
# Run in parallel and verify all results have the correct value
|
||||
results = run_functions_in_parallel(function_calls)
|
||||
|
||||
for _result_id, value in results.items():
|
||||
for value in results.values():
|
||||
assert value == "parallel_test"
|
||||
|
||||
|
||||
|
||||
@@ -131,9 +131,11 @@ def stream_chunks(
|
||||
chunks = []
|
||||
with client.stream("POST", "/v1/chat/completions", json=body) as response:
|
||||
assert response.status_code == 200
|
||||
for line in response.iter_lines():
|
||||
if line.startswith("data: ") and line != "data: [DONE]":
|
||||
chunks.append(json.loads(line[len("data: ") :]))
|
||||
chunks.extend(
|
||||
json.loads(line[len("data: ") :])
|
||||
for line in response.iter_lines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
|
||||
+1
-1
@@ -292,7 +292,7 @@ ignore = [
|
||||
"B904", # raise-without-from-inside-except (~465 existing violations).
|
||||
]
|
||||
# G004: f-strings in logging break Sentry's message-pattern grouping.
|
||||
select = ["ARG", "B", "E", "F", "G004", "I", "S", "W"]
|
||||
select = ["ARG", "B", "E", "F", "G004", "I", "PERF", "S", "W"]
|
||||
|
||||
[tool.ruff.lint.flake8-bugbear]
|
||||
# FastAPI dependency-injection markers are the intended way to write these
|
||||
|
||||
Reference in new issue
Block a user