vynte-infra/plane
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[PAI-957][PAI-875] feat: implemented ask build switcher (#4659)

Browse Source 
* feat: implemented ask build switcher

* fix: cleanup

* ask-build implementation in pi backend

* restrict reference to UUIDs in user facing response

* fix oauth flow in the build mode

* fix: OAuth message was being yielded twice

* add mode to chat preferences table

* clean generic_query_tool

* fix: stored the mode in the store

* pseudo streaming of final answer

* remove the no_actions_planned exit in build mode so that it can hanleretrieval only questions also

* migrate to sdk v0.2.0

* add json streaming in reasoning block

* clean unused tools from kit

* merge preview

* default mode for apps call to get-answer is build

* merge changes in Pi folder in preview branch

* fix reasoning header issues in ask mode

* clean no-actions-planned blocks

* logs for entity url tracking

* trim actions prompt. clean pk args from sdk methods

* fix: imlemented improved reasoning block

* project feature toggle before creation in planning mode

* fix nonetype error when entity urls is empty

* fix: reimplemented ask build changes

* fix: ticker

* fix: css

* fix table rendering issue in streaming api

* ask mode rejecting create table request as action

* remove old reasoning chunk stream to avoid duplications

* change header message for tool_message reasoning chunk

* remove results content from thinking block

* remove search active cycle result restriction to 1 entity

* add mode in rewrite. remove long reasoning instruction during tool calls loop

* retain project context during rewrite

* fix: implemneted chevron + fixed max height

* store source in msgs table

* add gpt-5.1 support

* fix rewrite failure due to chatrequest reconstruction error

* fix vector search result display in reasoning

* fix: reasoning header width fixed

* fix place holder resolution in execute-action

* fix project_id placeholder issue

* fix worklog sdk adapter bugs

* aggregate token counting for each step type

* clean action_planner.py code

* upgrade openai and langchain libs

* cleaning workitem tools

* fix: css

* fix: css + updated openai version

* clean tool args handling for workitems, projects, cycles and modules

* clean batch execution code

* feat: update batch execution and build mode helpers

* update links attachments and workitem types and properties tools

* beautify reasoning headers

* update issue_id to workitem_id in worklogs tools

* add pydantic models for types, comments and intake

* clean tool execution steps

* feat: restore apps/pi with latest changes after selective reset

- Restored all apps/pi changes including batch execution cleanup
- Deleted batch_action_orchestrator, batch_execution, and batch_execution_helpers
- Tool refinements and execution helpers updates
- SDK adapter improvements
- Reset other folders to cc5c927ba8 to avoid merge conflicts

* PAI - Fix the tool definitions of get-members api methods

* PAI fix still working msg formatting for mobile endpoint

* fix: build

* clean batch_execution code

* fix: refactor

* updade actions api/v2 endpoint/

* deduplicate tool names before sending to llm

* fix json parsing for claude models

* changes to action execution logic - exclude LLM intelligence

* clean build mode prompts

* add delete workitem tool

---------

Co-authored-by: sunder <[email protected]>
Co-authored-by: sunder <[email protected]>
This commit is contained in:
Akshita Goyal
2025-12-05 15:11:42 +05:30
committed by GitHub
co-authored by sunder sunder
parent 2740d73003
commit fc187a883f
98 changed files with 11736 additions and 12428 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/pi/docker-compose.yml
-10
View File
@@ -61,15 +61,6 @@ x-app-env: &app-env
AWS_S3_REGION: ${AWS_S3_REGION}
DEEPGRAM_API_KEY: ${DEEPGRAM_API_KEY}
x-feeder-env: &feeder-env
DEV_WORKSPACE_ID: ${DEV_WORKSPACE_ID}
FEED_ISSUES_DATA: ${FEED_ISSUES_DATA}
FEED_PAGES_DATA: ${FEED_PAGES_DATA}
FEED_DOCS_DATA: ${FEED_DOCS_DATA}
FEED_SLICES: ${FEED_SLICES}
BATCH_SIZE: ${BATCH_SIZE}
SCROLL_TIMEOUT: ${SCROLL_TIMEOUT}
x-celery-env: &celery-env
CELERY_BROKER_URL: ${CELERY_BROKER_URL:-pyamqp://${RABBITMQ_USER:-admin}:${RABBITMQ_PASSWORD:-admin123}@rabbitmq:5672//}
CELERY_RESULT_BACKEND: ${CELERY_RESULT_BACKEND:-rpc://}
@@ -79,7 +70,6 @@ x-celery-env: &celery-env
CELERY_VECTOR_SYNC_RETRY_DELAY: ${CELERY_VECTOR_SYNC_RETRY_DELAY:-60}
CELERY_WORKSPACE_PLAN_SYNC_ENABLED: ${CELERY_WORKSPACE_PLAN_SYNC_ENABLED:-1}
CELERY_WORKSPACE_PLAN_SYNC_INTERVAL: ${CELERY_WORKSPACE_PLAN_SYNC_INTERVAL:-86400}
DEV_WORKSPACE_ID: ${DEV_WORKSPACE_ID}
PLANE_PI_INTERNAL_API_URL: ${PLANE_PI_INTERNAL_API_URL}
PLANE_PI_INTERNAL_API_SECRET: ${PLANE_PI_INTERNAL_API_SECRET}
PLANE_PI_DATABASE_URL: ${PLANE_PI_DATABASE_URL}
apps/pi/pi/agents/sql_agent/base.py
+144 -14
View File
@@ -1,4 +1,5 @@
import json
import re
import time
from importlib.resources import read_text
from typing import Any
@@ -8,9 +9,10 @@ from typing import Union
from typing import cast
from uuid import UUID
from langchain.schema import AIMessage
from langchain.schema import HumanMessage
from langchain.schema import SystemMessage
from langchain_core.messages import AIMessage
from langchain_core.messages import HumanMessage
from langchain_core.messages import SystemMessage
from langchain_core.utils.json import parse_json_markdown
from pydantic import BaseModel
from rich.console import Console
from rich.panel import Panel
@@ -20,8 +22,10 @@ from sqlglot import parse_one
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import MessageMetaStepType
from pi.config import settings
from pi.services.chat.helpers.flow_tracking import FlowStepCollector
from pi.services.llm.error_handling import llm_error_handler
from pi.services.llm.llms import get_sql_agent_llm
from pi.services.schemas.chat import QueryFlowStore
@@ -129,7 +133,8 @@ async def _perform_table_selection_llm_call(
Note: For GPT-5 table selection, automatically uses gpt-5-fast instead of gpt-5-standard
to prevent token limit issues with large schema context. SQL generation keeps original model.
"""
# Create structured model dynamically and set tracking context
# Create model instance dynamically and set tracking context
# Use the provided model or fall back to the global one
if llm_model:
# For GPT-5 table selection, automatically reduce to fast variant to prevent token limits
@@ -143,9 +148,27 @@ async def _perform_table_selection_llm_call(
else:
table_selection_model_instance = table_selection_model
structured_table_selection_model = table_selection_model_instance.with_structured_output(TableSelectionResponse, include_raw=True) # type: ignore[arg-type]
structured_table_selection_model.set_tracking_context(message_id, db, MessageMetaStepType.SQL_TABLE_SELECTION, chat_id=chat_id) # type: ignore[attr-defined]
return await structured_table_selection_model.ainvoke(langchain_messages)
# Set tracking context on the model
table_selection_model_instance.set_tracking_context(message_id, db, MessageMetaStepType.SQL_TABLE_SELECTION, chat_id=chat_id) # type: ignore[attr-defined]
# Get raw response from LLM
raw_response = await table_selection_model_instance.ainvoke(langchain_messages)
# Extract content and parse JSON (handles markdown-wrapped JSON from Claude)
content = raw_response.content if hasattr(raw_response, "content") else str(raw_response)
try:
# parse_json_markdown handles both raw JSON and markdown-wrapped JSON
parsed_data = parse_json_markdown(content)
# Validate against Pydantic model
parsed_obj = TableSelectionResponse.model_validate(parsed_data)
return {"parsed": parsed_obj, "raw": raw_response}
except Exception as e:
log.error(f"ChatID: {chat_id} - Failed to parse table selection response: {e}")
log.error(f"ChatID: {chat_id} - Raw content: {content}")
# Fallback to empty table list
parsed_obj = TableSelectionResponse(relevant_tables=[])
return {"parsed": parsed_obj, "raw": raw_response}
# Function to select relevant tables for SQL query generation
@@ -258,11 +281,12 @@ async def sql_generation(
llm_model: Model to use for SQL generation
Returns:
Generated SQL query as string
Generated SQL query as string (with markdown wrappers stripped)
Note: For GPT-5, table selection uses gpt-5-fast to prevent token limits,
but SQL generation uses the original model for maximum quality.
"""
# Prepare messages for the LLM
langchain_messages: List[Message] = [SystemMessage(content=modified_sql_generator)]
for msg in messages:
@@ -277,12 +301,25 @@ async def sql_generation(
log.error("SQL generation failed after all retries")
return "SELECT 'Error: Unable to generate SQL query due to processing limitations' as error_message;"
# Ensure we always return a string - fixed type handling
# Extract content
if hasattr(response, "content"):
return str(response.content)
raw_sql = str(response.content)
else:
# Handle any other case by converting to string
return str(response)
raw_sql = str(response)
# Strip markdown code blocks (```sql...``` or ```...```)
# Use regex to detect if content is wrapped in code blocks
pattern = r"^```(?:sql)?\s*\n(.*?)\n```\s*$"
match = re.match(pattern, raw_sql.strip(), re.DOTALL)
if match:
# Extract just the SQL from markdown wrapper
clean_sql = match.group(1).strip()
else:
# No markdown wrapper, return as-is
clean_sql = raw_sql.strip()
return clean_sql
# SQL generation function
@@ -306,6 +343,14 @@ async def text2sql(
attachment_blocks: list[dict[str, Any]] | None = None,
) -> tuple[Dict[str, Any], Dict[str, Any]]:
try:
# Use caller-provided step_order from query_flow_store; no DB lookup
base_step_raw = query_flow_store.get("step_order")
try:
base_step = int(base_step_raw) if base_step_raw is not None else 0
except Exception:
base_step = 0
log.info(f"ChatID: {chat_id} - text2sql base_step from flow_store: {base_step}")
# Initialize a single dictionary to collect all intermediate results
intermediate_results: Dict[str, Any] = {
"steps": [], # Initialize steps as an empty list to avoid KeyError later
@@ -313,6 +358,9 @@ async def text2sql(
"entity_urls": None,
}
# Initialize flow step collector for text2sql sub-steps
flow_collector = FlowStepCollector(query_id=message_id, chat_id=str(chat_id) if chat_id else "", db=db) if message_id and chat_id else None
# Create user message for both preset and regular flows
user_message = _create_message_with_attachments(query, attachment_blocks)
@@ -433,6 +481,21 @@ async def text2sql(
# Store table validation step in our intermediate results
intermediate_results["post_processed_relevant_tables"] = relevant_tables
# Add table selection step to collector
if flow_collector:
flow_collector.add_step(
step_order=base_step + 1,
step_type=FlowStepType.TOOL,
tool_name="structured_db_tool_table_selection",
content={"relevant_tables": relevant_tables},
execution_data={
"relevant_tables": intermediate_results.get("relevant_tables"),
"post_processed_relevant_tables": relevant_tables,
},
is_planned=False,
is_executed=True,
)
# Step 2: Fetching whole schema for all the relevant tables
try:
relevant_tables_schemas = get_table_schemas(relevant_tables)
@@ -518,6 +581,22 @@ async def text2sql(
# Store SQL generation in our intermediate results
intermediate_results["generated_sql"] = generated_query
# Add preset SQL generation step to collector
if flow_collector:
flow_collector.add_step(
step_order=base_step + 2,
step_type=FlowStepType.TOOL,
tool_name="structured_db_tool_sql_generation",
content={"generated_sql": generated_query[:500]},
execution_data={
"generated_sql": generated_query,
"is_preset": True,
"relevant_tables": intermediate_results.get("post_processed_relevant_tables"),
},
is_planned=False,
is_executed=True,
)
else:
# Regular SQL generation with LLM
try:
@@ -570,6 +649,21 @@ async def text2sql(
# Store SQL generation in our intermediate results
intermediate_results["generated_sql"] = generated_query
# Add SQL generation step to collector
if flow_collector:
flow_collector.add_step(
step_order=base_step + 2,
step_type=FlowStepType.TOOL,
tool_name="structured_db_tool_sql_generation",
content={"generated_sql": generated_query[:500]},
execution_data={
"generated_sql": generated_query,
"relevant_tables": intermediate_results.get("post_processed_relevant_tables"),
},
is_planned=False,
is_executed=True,
)
except Exception as e:
log.error(f"Error during SQL generation for chat ID {chat_id}: {e}")
intermediate_results["sql_generation_error"] = e
@@ -602,8 +696,8 @@ async def text2sql(
generated_query_tables = set()
# Extract tables from all SELECT statements in the query
for select in parsed_query.find_all(exp.Select):
select_tables = _get_available_tables(select)
for select_node in parsed_query.find_all(exp.Select):
select_tables = _get_available_tables(select_node)
generated_query_tables.update(select_tables)
# Find extra tables that need to be added to CTE
@@ -633,6 +727,21 @@ async def text2sql(
# Store CTE generation in our intermediate results
intermediate_results["cte_head"] = CTE_head
# Add CTE generation step to collector
if flow_collector:
flow_collector.add_step(
step_order=base_step + 3,
step_type=FlowStepType.TOOL,
tool_name="structured_db_tool_cte_generation",
content={"cte_head": CTE_head[:500]},
execution_data={
"cte_head": CTE_head,
"tables_for_cte": tables_for_cte,
},
is_planned=False,
is_executed=True,
)
except Exception as e:
log.error(f"Error generating CTE query for chat ID {chat_id}: {e}")
intermediate_results["cte_generation_error"] = e
@@ -658,6 +767,22 @@ async def text2sql(
intermediate_results["final_query"] = final_query
# Add final query step to collector
if flow_collector:
flow_collector.add_step(
step_order=base_step + 4,
step_type=FlowStepType.TOOL,
tool_name="structured_db_tool_final_query",
content={"final_query": final_query[:500]},
execution_data={
"final_query": final_query,
"generated_sql": intermediate_results.get("generated_sql"),
"cte_head": intermediate_results.get("cte_head"),
},
is_planned=False,
is_executed=True,
)
# Step 6: SQL Execution
query_execution_result: Any
try:
@@ -757,6 +882,11 @@ async def text2sql(
response_data: Dict[str, Any] = {"sql_query": final_query, "results": query_execution_result}
# Persist all collected flow steps
if flow_collector:
async with flow_collector:
pass # Steps are persisted in __aexit__
return intermediate_results, response_data
except Exception as e:
apps/pi/pi/agents/sql_agent/prompts.py
+24 -44
View File
@@ -168,57 +168,37 @@ Here is the list of available tables in Plane's database with their descriptions
{table_descriptions}
</table_descriptions>
Your goal is to select only the tables that are directly relevant to fulfilling this query. Follow these steps:
Your goal is to select only the tables that are directly relevant to fulfilling this query.
1. Analyze the user's query to identify key elements and requirements.
2. Review each table in the table_descriptions, considering:
- The table's purpose and description
- What the table contains and does not contain
- The table's relationships with other tables
3. Categorize each table as "Highly Relevant", "Possibly Relevant", or "Not Relevant" to the query.
4. Select the tables that are essential for answering the query.
5. If a table contains the data directly related to the query, classify it as 'Highly Relevant' immediately.
6. Format your final response as a JSON object with a "relevant_tables" key containing an array of selected table names.
**Selection Process:**
1. Identify key entities and concepts mentioned in the user query (e.g., issues, users, projects, cycles, states).
2. For each identified concept, find the corresponding table(s) that contain that data.
3. Include related tables needed for joins or filtering.
4. Include related tables needed to display human-readable names (e.g., if selecting `issue_assignees`, also include `users` table).
Before providing your final answer, wrap your analysis in <table_analysis> tags. In this analysis:
1. Summarize the key elements from the user query.
2. For each table in the table_descriptions:
a. Briefly describe its purpose
b. Explain its relevance (or lack thereof) to the query
c. Categorize it as "Highly Relevant", "Possibly Relevant", or "Not Relevant"
3. *Important Note*
Do not ignore tables whose descriptions directly match the query, even if they have relationships with other tables.
For example, for queries about "users," always consider the users table as Highly Relevant.
4. Justify your final selections based on this categorization
**Incomplete/Active Work:**
- When the user's intent implies items that should be worked on NOW (e.g., "to add to the
current cycle", "missed to add", "should be in this cycle", "to pick up", "pending",
"open", "to-do"), and the user has not explicitly asked for completed/closed items,
you MUST include the states table in the selection.
It's OK for this section to be quite long.
Example of the expected output format:
{{"relevant_tables": ["table1", "table2", "table3"]}}
Additional Instructions Related to Issue Ownership:
- Ownership of an issue lies with the users assigned to it. Not with the ones who created it.
- When determining ownership of an issue, always use the issue_assignees table as the primary source. Treat it as the most relevant table for any query about who an issue belongs to, such as ownership, assignments, or responsibility. Do not prioritize issues table alone for such cases.
**Issue Ownership:**
- Ownership of an issue lies with the users assigned to it, not with the ones who created it.
- When determining ownership, always include BOTH `issue_assignees` AND `users` tables.
- Phrases containing "my/mine" related to issues:
- "My issues" = Issues where the user is an assignee. Use issue_assignees for this.
- "My backlog" = Issues in backlog state where user is an assignee. Use issue_assignees for this.
- "Created by me" = Issues where created_by_id matches the user. Use issues table for this.
- When both creation and ownership appear ("my issues that I created"), use BOTH issue_assignees and issues.
- "My issues" → include `issue_assignees`, `users`, `states` tables
- "My backlog" → include `issue_assignees`, `users`, `states` tables
- "Created by me" → include `issues`, `users`, `states` tables
- "Issues assigned to [person]" → include `issue_assignees`, `users`, `states` tables
Additional Instructions Related to Incomplete/Active Work:
- For queries that imply items that should be worked on (e.g., "to add to the current cycle", "missed to add", "to pick up", "open", "pending", "not done"), prioritize uncompleted states.
- Treat the states table as Highly Relevant and include it in selections when the query does not specify states but implies active/to-do/pending items.
- Default uncompleted buckets/groups: backlog, unstarted, started. Exclude completed/cancelled unless explicitly requested by the user.
**Priority Queries:**
- For priority-related queries, include the `issues` table (which contains the priority attribute).
Additional Instructions Related to Priority Canonicalization:
- Canonical priority values are: urgent, high, medium, low (case-insensitive). No other priority names should be assumed.
- Map common synonyms to canonical values for interpretation: "highest"/"critical"/"blocker"/"p0" → urgent; "very high"/"p1" → high; "normal"/"standard" → medium; "lowest" → low.
- Consider the priority attribute on the issues table the canonical source when selecting tables for priority-driven queries.
Important requirements:
**Output Format:**
- Your response must be ONLY a valid JSON object with this exact format:
{{"relevant_tables": ["table1", "table2", "table3"]}}
- Include ONLY tables that exist in the provided table_descriptions.
- Select only tables that are directly relevant to the query or necessary due to relationships.
- Do not include any additional text, explanations, or analysis outside the JSON object.
- Do not generate SQL queries or any other extraneous information.
- Ensure your final output is a valid JSON object.
- Do not include any additional text or explanations outside the JSON object.
"""
apps/pi/pi/agents/sql_agent/tools.py
+9 -9
View File
@@ -680,7 +680,7 @@ def extract_ids_from_sql_result(query_result: Any) -> Dict[str, List[str]]:
async def get_refs_from_chat(db: AsyncSession, chat_id: str) -> List[str]:
"""Safely extract any available entity URLs from execution_data or content."""
tool_results = await get_tool_results_from_chat_history(db, uuid.UUID(chat_id), "PLANE_STRUCTURED_DATABASE_AGENT")
tool_results = await get_tool_results_from_chat_history(db, uuid.UUID(chat_id), "structured_db_tool")
extracted_entities: List[str] = []
@@ -716,7 +716,7 @@ async def get_refs_from_chat(db: AsyncSession, chat_id: str) -> List[str]:
async def construct_entity_urls_vectordb(entity_ids: Dict[str, List[str]], api_base_url: str) -> List[Dict[str, str]]:
"""Construct entity links using OpenSearch indices for better efficiency"""
entity_links: List[Dict[str, str]] = []
log.info(f"construct_entity_urls_vectordb: Entity IDs: {entity_ids}")
if not entity_ids or not any(entity_ids.values()):
return entity_links
@@ -743,7 +743,7 @@ async def construct_entity_urls_vectordb(entity_ids: Dict[str, List[str]], api_b
try:
# Build batch query to get all documents by IDs
query = {"query": {"terms": {"id": entity_list}}, "size": len(entity_list)}
log.info(f"construct_entity_urls_vectordb: Query: {query}")
# Execute search
response = await vectorstore.async_search(index=index_name, body=query)
@@ -751,11 +751,11 @@ async def construct_entity_urls_vectordb(entity_ids: Dict[str, List[str]], api_b
for hit in response["hits"]["hits"]:
source = hit["_source"]
entity_id = source.get("id", "")
log.info(f"construct_entity_urls_vectordb: Source: {source}")
# Extract common fields
name = source.get("name", "")
workspace_slug = source.get("workspace_slug", "")
log.info(f"construct_entity_urls_vectordb: Name: {name}, Workspace Slug: {workspace_slug}")
# Construct URLs based on entity type
if entity_type == "issues":
# For issues: /workspace_slug/browse/PROJECT_IDENTIFIER-SEQUENCE_ID/
@@ -766,7 +766,7 @@ async def construct_entity_urls_vectordb(entity_ids: Dict[str, List[str]], api_b
url = f"{api_base_url}/{workspace_slug}/browse/{issue_identifier}/"
entity_link = {"name": name, "id": entity_id, "issue_identifier": issue_identifier, "url": url, "type": "issue"}
entity_links.append(entity_link)
log.info(f"construct_entity_urls_vectordb: Entity (workitem) Link: {entity_link}")
elif entity_type == "pages":
# For pages: global vs project pages
is_global = source.get("is_global", False)
@@ -784,13 +784,13 @@ async def construct_entity_urls_vectordb(entity_ids: Dict[str, List[str]], api_b
entity_link = {"name": name, "id": entity_id, "url": url, "type": "page"}
entity_links.append(entity_link)
log.info(f"construct_entity_urls_vectordb: Entity (page) Link: {entity_link}")
elif entity_type == "projects":
# For projects: /workspace_slug/projects/project_id/overview/
url = f"{api_base_url}/{workspace_slug}/projects/{entity_id}/overview/"
entity_link = {"name": name, "id": entity_id, "url": url, "type": "project"}
entity_links.append(entity_link)
log.info(f"construct_entity_urls_vectordb: Entity (project) Link: {entity_link}")
elif entity_type in ["modules", "cycles"]:
# For modules and cycles: /workspace_slug/projects/project_id/modules|cycles/entity_id/
project_id = source.get("project_id", "")
@@ -804,7 +804,7 @@ async def construct_entity_urls_vectordb(entity_ids: Dict[str, List[str]], api_b
"type": entity_type[:-1], # Remove 's' from end (modules -> module)
}
entity_links.append(entity_link)
log.info(f"construct_entity_urls_vectordb: Entity (module or cycle) Link: {entity_link}")
except Exception as e:
log.error(f"Error constructing URLs for {entity_type}: {e}")
continue
apps/pi/pi/alembic/versions/74852ee169e1_polymorphic_focus_context.py
+76
View File
@@ -0,0 +1,76 @@
"""polymorphic focus context
Revision ID: 74852ee169e1
Revises: 6082d91ecc4e
Create Date: 2025-01-28
This migration adds polymorphic focus context columns (focus_entity_type, focus_entity_id)
to replace the specific focus_project_id and focus_workspace_id columns.
This allows supporting multiple entity types (workspace, project, cycle, module, etc.)
without adding new columns for each type.
Also adds mode column to store chat mode (ask/build).
Also adds source column to messages table to track the source of chat requests.
"""
from typing import Sequence
from typing import Union
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "74852ee169e1"
down_revision: Union[str, None] = "6082d91ecc4e"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# Add new polymorphic columns
op.add_column("user_chat_preferences", sa.Column("focus_entity_type", sa.String(length=50), nullable=True))
op.add_column("user_chat_preferences", sa.Column("focus_entity_id", sa.Uuid(), nullable=True))
# Add mode column with default value 'ask'
op.add_column("user_chat_preferences", sa.Column("mode", sa.String(length=10), nullable=False, server_default="ask"))
# Add source column to messages table to track the source of chat requests
op.add_column("messages", sa.Column("source", sa.String(length=50), nullable=True))
# Migrate existing data from old columns to new polymorphic structure
# If focus_project_id is set, migrate to focus_entity_type='project' and focus_entity_id=focus_project_id
op.execute("""
UPDATE user_chat_preferences
SET focus_entity_type = 'project', focus_entity_id = focus_project_id
WHERE focus_project_id IS NOT NULL
""")
# If focus_workspace_id is set (and focus_project_id is NULL), migrate to focus_entity_type='workspace' and focus_entity_id=focus_workspace_id
op.execute("""
UPDATE user_chat_preferences
SET focus_entity_type = 'workspace', focus_entity_id = focus_workspace_id
WHERE focus_workspace_id IS NOT NULL AND focus_project_id IS NULL
""")
# Note: We keep the old columns (focus_project_id, focus_workspace_id) for backward compatibility
# They will be dropped in a future migration after all code is updated
def downgrade() -> None:
# Migrate data back from new columns to old columns
op.execute("""
UPDATE user_chat_preferences
SET focus_project_id = focus_entity_id
WHERE focus_entity_type = 'project' AND focus_entity_id IS NOT NULL
""")
op.execute("""
UPDATE user_chat_preferences
SET focus_workspace_id = focus_entity_id
WHERE focus_entity_type = 'workspace' AND focus_entity_id IS NOT NULL
""")
# Drop new columns
op.drop_column("messages", "source")
op.drop_column("user_chat_preferences", "mode")
op.drop_column("user_chat_preferences", "focus_entity_id")
op.drop_column("user_chat_preferences", "focus_entity_type")
apps/pi/pi/app/api/v1/endpoints/_sse.py
+4 -1
View File
@@ -17,7 +17,7 @@ def sse_done() -> str:
return sse_event("done", {"done": True})
def normalize_error_chunk(chunk: str) -> Optional[str]:
def normalize_error_chunk(chunk: Any) -> Optional[str]:
"""Convert plain text error chunks to a proper SSE error event.
Recognizes patterns like "error: <message>" (with optional quotes) and
@@ -25,6 +25,9 @@ def normalize_error_chunk(chunk: str) -> Optional[str]:
Returns a formatted SSE error event string or None if not applicable.
"""
if not isinstance(chunk, str):
return None
if chunk.startswith("error:"):
msg = chunk.split("error:", 1)[1].strip()
if (msg.startswith("'") and msg.endswith("'")) or (msg.startswith('"') and msg.endswith('"')):
apps/pi/pi/app/api/v1/endpoints/chat.py
+99 -102
View File
@@ -4,7 +4,9 @@ import json
from typing import Any
from typing import AsyncGenerator
from typing import Coroutine
from typing import Dict
from typing import Optional
from typing import Union
from typing import cast
from urllib.parse import urlparse
from uuid import UUID
@@ -28,11 +30,7 @@ from pi.app.api.v1.dependencies import validate_plane_token
from pi.app.api.v1.endpoints._sse import normalize_error_chunk
from pi.app.api.v1.endpoints._sse import sse_done
from pi.app.api.v1.endpoints._sse import sse_event
from pi.app.api.v1.helpers.batch_execution_helpers import execute_batch_actions
from pi.app.api.v1.helpers.batch_execution_helpers import format_execution_response
from pi.app.api.v1.helpers.batch_execution_helpers import prepare_execution_data
from pi.app.api.v1.helpers.batch_execution_helpers import update_assistant_message_with_execution_results
from pi.app.api.v1.helpers.batch_execution_helpers import validate_session_and_get_user
from pi.app.api.v1.helpers.execution import chosen_llm
from pi.app.api.v1.helpers.plane_sql_queries import resolve_workspace_id_from_project_id
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import UserTypeChoices
@@ -61,9 +59,8 @@ from pi.app.utils.exceptions import SQLGenerationError
from pi.core.db.plane_pi.lifecycle import get_async_session
from pi.core.db.plane_pi.lifecycle import get_streaming_db_session
from pi.services.actions.oauth_service import PlaneOAuthService
from pi.services.chat.action_executor import BuildModeToolExecutor
from pi.services.chat.chat import PlaneChatBot
from pi.services.chat.helpers.batch_execution_helpers import get_original_user_query
from pi.services.chat.helpers.batch_execution_helpers import get_planned_actions_for_execution
from pi.services.chat.helpers.tool_utils import format_clarification_as_text
from pi.services.chat.search import ChatSearchService
from pi.services.chat.templates import tiles_factory
@@ -183,17 +180,19 @@ async def get_answer_for_slack(data: ChatRequest, request: Request, db: AsyncSes
final_response = ""
actions_data = {}
clarification_data = {}
formatted_context = {}
formatted_context: dict[str, Any] = {}
response_type = "response"
# listen to all the stream chunks, join the chunks and return the complete response
# as json object
async with get_streaming_db_session() as stream_db:
base_iter = chatbot.process_query_stream(data, db=stream_db)
data.mode = "build" # type: ignore
base_iter = chatbot.process_chat_stream(data, db=stream_db)
async for chunk in base_iter:
# Ignore all intermediate chunks
if chunk.startswith("πspecial reasoning block"):
if isinstance(chunk, dict):
# Currently only reasoning chunk is sent as dict.
continue
# Ignore all intermediate chunks
if chunk.startswith("πspecial actions blockπ: "):
response_type = "actions"
actions_data = json.loads(chunk.replace("πspecial actions blockπ: ", ""))
@@ -232,7 +231,9 @@ async def get_answer_for_slack(data: ChatRequest, request: Request, db: AsyncSes
else:
return JSONResponse(status_code=400, content={"detail": "Invalid artifact_id format"})
execution_data = await prepare_execution_data(
# Execute batch actions using the service
service = BuildModeToolExecutor(chatbot=PlaneChatBot("gpt-4.1"), db=db)
result = await service.execute(
ActionBatchExecutionRequest(
workspace_id=workspace_id,
chat_id=chat_id,
@@ -241,29 +242,23 @@ async def get_answer_for_slack(data: ChatRequest, request: Request, db: AsyncSes
access_token=access_token,
),
user_id,
db,
)
if not execution_data:
# Determine the specific error based on what failed
if not await get_planned_actions_for_execution(message_id, chat_id, db):
return JSONResponse(status_code=404, content={"detail": BATCH_EXECUTION_ERRORS["NO_PLANNED_ACTIONS"]})
elif not await get_original_user_query(message_id, db):
return JSONResponse(status_code=404, content={"detail": BATCH_EXECUTION_ERRORS["NO_ORIGINAL_QUERY"]})
else:
# OAuth or workspace issue
return JSONResponse(
status_code=401,
content={
"detail": BATCH_EXECUTION_ERRORS["OAUTH_REQUIRED"],
"error_code": "OAUTH_REQUIRED",
"workspace_id": str(data.workspace_id),
"user_id": str(data.user_id),
},
)
# Execute batch actions
context = await execute_batch_actions(execution_data, db)
formatted_context = format_execution_response(context)
# Check if service returned an error
if result.get("error"):
status_code = result.get("status_code", 500)
detail = result.get("detail", "Unknown error")
content = {"detail": detail}
if "error_code" in result:
content["error_code"] = result["error_code"]
if "workspace_id" in result:
content["workspace_id"] = result["workspace_id"]
if "user_id" in result:
content["user_id"] = result["user_id"]
return JSONResponse(status_code=status_code, content=content)
# Extract the formatted context from successful result
formatted_context = result
return JSONResponse(
status_code=200,
@@ -300,7 +295,6 @@ async def get_answer(data: ChatRequest, session: str = Depends(cookie_schema)):
async def stream_response() -> AsyncGenerator[str, None]:
token_id = None
try:
pending_backticks = ""
# Open a short-lived session for the duration of the streaming work
async with get_streaming_db_session() as stream_db:
# Heartbeat mechanism that does not cancel the underlying generator
@@ -313,8 +307,11 @@ async def get_answer(data: ChatRequest, session: str = Depends(cookie_schema)):
with contextlib.suppress(asyncio.CancelledError):
await asyncio.sleep(10)
base_iter = chatbot.process_query_stream(data, db=stream_db)
next_chunk_task: asyncio.Task[str] = asyncio.create_task(cast(Coroutine[None, None, str], base_iter.__anext__()))
# Single unified call; internal routing happens in chat service
base_iter = chatbot.process_chat_stream(data, db=stream_db)
next_chunk_task: asyncio.Task[Union[str, Dict[str, Any]]] = asyncio.create_task(
cast(Coroutine[None, None, Union[str, Dict[str, Any]]], base_iter.__anext__())
)
# Start initial heartbeat timer
heartbeat_task = asyncio.create_task(heartbeat_emitter())
@@ -326,7 +323,7 @@ async def get_answer(data: ChatRequest, session: str = Depends(cookie_schema)):
# If heartbeat timer completed first, emit heartbeat and restart timer
if heartbeat_task in done and not heartbeat_stop.is_set():
payload = {"reasoning": "⏳ Still working...\n\n"}
payload: Dict[str, Any] = {"header": "⏳ Still working...\n\n", "content": ""}
yield f"event: reasoning\ndata: {json.dumps(payload)}\n\n"
# Restart heartbeat timer
heartbeat_task = asyncio.create_task(heartbeat_emitter())
@@ -348,17 +345,26 @@ async def get_answer(data: ChatRequest, session: str = Depends(cookie_schema)):
heartbeat_task = asyncio.create_task(heartbeat_emitter())
# Normalize plain-text error chunks to SSE error events
normalized_error = normalize_error_chunk(chunk)
normalized_error = normalize_error_chunk(chunk) if isinstance(chunk, str) else None
if normalized_error:
yield normalized_error
next_chunk_task = asyncio.create_task(cast(Coroutine[None, None, str], base_iter.__anext__()))
next_chunk_task = asyncio.create_task(cast(Coroutine[None, None, Union[str, Dict[str, Any]]], base_iter.__anext__()))
continue
if chunk.startswith("πspecial reasoning blockπ: "):
reasoning_content = chunk.replace("πspecial reasoning blockπ: ", "")
payload = {"reasoning": reasoning_content}
yield f"event: reasoning\ndata: {json.dumps(payload)}\n\n"
elif chunk.startswith("πspecial actions blockπ: "):
# change point: "πspecial reasoning blockπ:
# handle the new json string format of reasoning blocks. It is a json string with header and content.
# add a condition to check if the chunk is a json string with header and content.
if isinstance(chunk, dict):
if "chunk_type" in chunk and chunk["chunk_type"] == "reasoning":
payload = {"header": chunk["header"], "content": chunk["content"]}
yield f"event: reasoning\ndata: {json.dumps(payload)}\n\n"
# for backward compatibility, we also need to yield the chunk as a string
# bwc_payload = {"reasoning": f"{chunk["header"]}\n\n{chunk["content"]}\n\n"}
# yield f"event: reasoning\ndata: {json.dumps(bwc_payload)}\n\n"
else:
log.warning(f"ChatID: {data.chat_id} - Chunk is not a json string: {chunk}")
payload = {"chunk": chunk}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
elif isinstance(chunk, str) and chunk.startswith("πspecial actions blockπ: "):
actions_content = chunk.replace("πspecial actions blockπ: ", "")
try:
actions_data = json.loads(actions_content)
@@ -367,11 +373,13 @@ async def get_answer(data: ChatRequest, session: str = Depends(cookie_schema)):
log.warning(f"Failed to parse actions JSON: {actions_content}")
payload = {"chunk": chunk}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
elif chunk.startswith("πspecial clarification blockπ: "):
elif isinstance(chunk, str) and chunk.startswith("πspecial clarification blockπ: "):
clarification_content = chunk.replace("πspecial clarification blockπ: ", "")
try:
clarification_data = json.loads(clarification_content)
log.info(f"ChatID: {data.chat_id} - Clarification data received in the endpoint: {clarification_data}")
formatted_text = format_clarification_as_text(clarification_data)
log.info(f"ChatID: {data.chat_id} - Clarification formatted text: {formatted_text}")
payload = {"chunk": formatted_text}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
except json.JSONDecodeError:
@@ -381,24 +389,8 @@ async def get_answer(data: ChatRequest, session: str = Depends(cookie_schema)):
}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
else:
# Handle code block formatting
if chunk.startswith("```"):
if pending_backticks:
payload = {"chunk": pending_backticks}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
pending_backticks = ""
payload = {"chunk": chunk}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
elif chunk in ["`", "``"]:
pending_backticks += chunk
next_chunk_task = asyncio.create_task(cast(Coroutine[None, None, str], base_iter.__anext__()))
continue
else:
if pending_backticks:
chunk = pending_backticks + chunk
pending_backticks = ""
payload = {"chunk": chunk}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
payload = {"chunk": chunk}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
# Prepare next iteration
next_chunk_task = asyncio.create_task(cast(Coroutine[None, None, str], base_iter.__anext__()))
@@ -409,11 +401,6 @@ async def get_answer(data: ChatRequest, session: str = Depends(cookie_schema)):
with contextlib.suppress(asyncio.CancelledError):
await heartbeat_task
# Handle any remaining pending backticks at the end of stream
if pending_backticks:
payload = {"chunk": pending_backticks}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
# Extract token_id from data.context if available for background task
if hasattr(data, "context") and isinstance(data.context, dict):
token_id = data.context.get("token_id")
@@ -587,8 +574,11 @@ async def get_chat_history(
"feedback": results.get("feedback", ""),
"reasoning": results.get("reasoning", ""),
"is_focus_enabled": results.get("is_focus_enabled", False),
"focus_entity_type": results.get("focus_entity_type", None),
"focus_entity_id": results.get("focus_entity_id", None),
"focus_project_id": results.get("focus_project_id", None),
"focus_workspace_id": results.get("focus_workspace_id", None),
"mode": results.get("mode", "ask"),
},
},
)
@@ -604,8 +594,11 @@ async def get_chat_history(
"feedback": results.get("feedback", ""),
"reasoning": results.get("reasoning", ""),
"is_focus_enabled": results.get("is_focus_enabled", False),
"focus_entity_type": results.get("focus_entity_type", None),
"focus_entity_id": results.get("focus_entity_id", None),
"focus_project_id": results.get("focus_project_id", None),
"focus_workspace_id": results.get("focus_workspace_id", None),
"mode": results.get("mode", "ask"),
},
},
)
@@ -619,8 +612,11 @@ async def get_chat_history(
"feedback": results["feedback"],
"reasoning": results.get("reasoning", ""),
"is_focus_enabled": results.get("is_focus_enabled", False),
"focus_entity_type": results.get("focus_entity_type", None),
"focus_entity_id": results.get("focus_entity_id", None),
"focus_project_id": results.get("focus_project_id", None),
"focus_workspace_id": results.get("focus_workspace_id", None),
"mode": results.get("mode", "ask"),
}
}
)
@@ -671,8 +667,11 @@ async def get_chat_history_object(
"feedback": results["feedback"],
"reasoning": results.get("reasoning", ""),
"is_focus_enabled": results.get("is_focus_enabled", False),
"focus_entity_type": results.get("focus_entity_type", None),
"focus_entity_id": results.get("focus_entity_id", None),
"focus_project_id": results.get("focus_project_id", None),
"focus_workspace_id": results.get("focus_workspace_id", None),
"mode": results.get("mode", "ask"),
}
}
)
@@ -947,6 +946,7 @@ async def queue_answer(data: ChatRequest, db: AsyncSession = Depends(get_async_s
user_type=UserTypeChoices.USER.value,
llm_model=data.llm,
workspace_slug=resolved_workspace_slug,
source=data.source or None,
db=db,
)
@@ -1029,14 +1029,14 @@ async def stream_answer(token: UUID4, db: AsyncSession = Depends(get_async_sessi
if not flow_step:
# REGENERATE FLOW: No QUEUE flow step means this is regenerate
log.info(f"No QUEUE flow step found for token {token}. Treating as regenerate request.")
# 1. Get the user message
# Get the user message
user_message = await get_message_by_id(db, token)
if not user_message:
log.warning(f"Message not found for token {token}")
return JSONResponse(status_code=404, content={"detail": "Message not found"})
# 2. Mark old assistant response as replaced BEFORE generating new one
# Mark old assistant response as replaced BEFORE generating new one
# This ensures retrieve_chat_history() won't include it in context for new generation
marked = await mark_assistant_response_as_replaced(db, user_message.id)
if marked:
@@ -1044,14 +1044,14 @@ async def stream_answer(token: UUID4, db: AsyncSession = Depends(get_async_sessi
else:
log.info(f"No existing assistant response found for message {token} (first generation or already replaced)")
# 4. Reconstruct ChatRequest from user message
# Reconstruct ChatRequest from user message
try:
queued_request = await reconstruct_chat_request_from_message(db, user_message, user_id)
except Exception as e:
log.error(f"Error reconstructing ChatRequest from message {token}: {e}")
return JSONResponse(status_code=500, content={"detail": "Failed to reconstruct request"})
# 5. Pass token_id so new assistant message reuses same user message as parent
# Pass token_id so new assistant message reuses same user message as parent
try:
queued_request.context["token_id"] = str(token)
except Exception as e:
@@ -1059,7 +1059,7 @@ async def stream_answer(token: UUID4, db: AsyncSession = Depends(get_async_sessi
log.info(f"Regenerating response for message {token}")
# 6. Stream new response (get_answer will call process_query_stream)
# Stream new response (get_answer will call process_query_stream)
# When process_query_stream calls retrieve_chat_history,
# it will NOT see the old response because is_replaced=True
return await get_answer(data=queued_request, session=session)
@@ -1123,40 +1123,37 @@ async def execute_action(request: ActionBatchExecutionRequest, db: AsyncSession
# 3. Assistant messages are updated with execution results and entity information
# 4. Conversation history includes explicit text about executed vs. not executed actions
# This ensures complete context for follow-up questions and LLM understanding.
try:
# Validate session and get user
user_id = await validate_session_and_get_user(session)
if not user_id:
return JSONResponse(status_code=401, content={"detail": BATCH_EXECUTION_ERRORS["INVALID_SESSION"]})
auth = await is_valid_session(session)
if not auth.user:
return JSONResponse(status_code=401, content={"detail": "Invalid User"})
# Validate and prepare execution data
execution_data = await prepare_execution_data(request, user_id, db)
if not execution_data:
# Determine the specific error based on what failed
if not await get_planned_actions_for_execution(request.message_id, request.chat_id, db):
return JSONResponse(status_code=404, content={"detail": BATCH_EXECUTION_ERRORS["NO_PLANNED_ACTIONS"]})
elif not await get_original_user_query(request.message_id, db):
return JSONResponse(status_code=404, content={"detail": BATCH_EXECUTION_ERRORS["NO_ORIGINAL_QUERY"]})
else:
# OAuth or workspace issue
return JSONResponse(
status_code=401,
content={
"detail": BATCH_EXECUTION_ERRORS["OAUTH_REQUIRED"],
"error_code": "OAUTH_REQUIRED",
"workspace_id": str(request.workspace_id),
"user_id": str(user_id),
},
)
llm_model = await chosen_llm(db=db, message_id=request.message_id)
# Use default model if none was found in the message
chatbot = PlaneChatBot(llm_model or "gpt-4.1")
build_mode_tool_executor = BuildModeToolExecutor(chatbot=chatbot, db=db)
result = await build_mode_tool_executor.execute(request, auth.user.id)
# Execute batch actions
context = await execute_batch_actions(execution_data, db)
# Check if service returned an error
if result.get("error"):
status_code = result.get("status_code", 500)
detail = result.get("detail", "Unknown error")
# Update the assistant message with execution results
await update_assistant_message_with_execution_results(request.message_id, request.chat_id, context, db)
# Build response content
content = {"detail": detail}
if "error_code" in result:
content["error_code"] = result["error_code"]
if "workspace_id" in result:
content["workspace_id"] = result["workspace_id"]
if "user_id" in result:
content["user_id"] = result["user_id"]
# Return appropriate response based on execution status
return JSONResponse(content=format_execution_response(context))
return JSONResponse(status_code=status_code, content=content)
# Return successful response
return JSONResponse(content=result)
except Exception as e:
log.error(f"Error in execute_action: {str(e)}")
apps/pi/pi/app/api/v1/endpoints/mobile/chat.py
+32 -304
View File
@@ -3,8 +3,11 @@ import asyncio
import contextlib
import json
import uuid
from typing import Any
from typing import Coroutine
from typing import Dict
from typing import Optional
from typing import Union
from typing import cast
from fastapi import APIRouter
@@ -18,12 +21,7 @@ from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.api.v1.dependencies import jwt_schema
from pi.app.api.v1.dependencies import validate_jwt_token
from pi.app.api.v1.helpers.plane_sql_queries import resolve_workspace_id_from_project_id
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import UserTypeChoices
from pi.app.models.message import MessageFlowStep
from pi.app.schemas.mobile.chat import ChatFeedbackMobile
from pi.app.schemas.mobile.chat import ChatInitializationRequestMobile
from pi.app.schemas.mobile.chat import ChatRequestMobile
from pi.app.schemas.mobile.chat import ChatSearchResponseMobile
from pi.app.schemas.mobile.chat import ChatStartResponseMobile
@@ -36,8 +34,6 @@ from pi.app.schemas.mobile.chat import ModelsResponseMobile
from pi.app.schemas.mobile.chat import RenameChatRequestMobile
from pi.app.schemas.mobile.chat import TitleRequestMobile
from pi.app.schemas.mobile.chat import UnfavoriteChatRequestMobile
from pi.app.utils import validate_chat_initialization
from pi.app.utils import validate_chat_request
from pi.app.utils.background_tasks import schedule_chat_deletion
from pi.app.utils.background_tasks import schedule_chat_rename
from pi.app.utils.background_tasks import schedule_chat_search_upsert
@@ -48,8 +44,6 @@ from pi.services.chat.chat import PlaneChatBot
from pi.services.chat.helpers.tool_utils import format_clarification_as_text
from pi.services.chat.search import ChatSearchService
from pi.services.chat.templates import tiles_factory
from pi.services.chat.utils import initialize_new_chat
from pi.services.chat.utils import resolve_workspace_slug
from pi.services.retrievers.pg_store import favorite_chat
from pi.services.retrievers.pg_store import get_active_models
from pi.services.retrievers.pg_store import get_chat_messages
@@ -60,11 +54,6 @@ from pi.services.retrievers.pg_store import retrieve_chat_history
from pi.services.retrievers.pg_store import soft_delete_chat
from pi.services.retrievers.pg_store import unfavorite_chat
from pi.services.retrievers.pg_store import update_message_feedback
from pi.services.retrievers.pg_store.message import get_message_by_id
from pi.services.retrievers.pg_store.message import mark_assistant_response_as_replaced
from pi.services.retrievers.pg_store.message import reconstruct_chat_request_from_message
from pi.services.retrievers.pg_store.message import upsert_message
from pi.services.retrievers.pg_store.message import upsert_message_flow_steps
log = logger.getChild("v1/mobile/chat")
mobile_router = APIRouter()
@@ -109,8 +98,10 @@ async def get_answer(
with contextlib.suppress(asyncio.CancelledError):
await asyncio.sleep(10)
base_iter = chatbot.process_query_stream(data, stream_db)
next_chunk_task: asyncio.Task[str] = asyncio.create_task(cast(Coroutine[None, None, str], base_iter.__anext__()))
base_iter = chatbot.process_chat_stream(data, stream_db)
next_chunk_task: asyncio.Task[Union[str, Dict[str, Any]]] = asyncio.create_task(
cast(Coroutine[None, None, Union[str, Dict[str, Any]]], base_iter.__anext__())
)
# Start initial heartbeat timer
heartbeat_task = asyncio.create_task(heartbeat_emitter())
@@ -122,7 +113,8 @@ async def get_answer(
# If heartbeat timer completed first, emit heartbeat and restart timer
if heartbeat_task in done and not heartbeat_stop.is_set():
yield "event: reasoning\ndata: ⏳ Still working...\n\n"
heartbeat_payload = {"reasoning": "⏳ Still working..."}
yield f"event: reasoning\ndata: {json.dumps(heartbeat_payload)}\n\n"
# Restart heartbeat timer
heartbeat_task = asyncio.create_task(heartbeat_emitter())
@@ -142,10 +134,17 @@ async def get_answer(
await heartbeat_task
heartbeat_task = asyncio.create_task(heartbeat_emitter())
if chunk.startswith("πspecial reasoning blockπ: "):
reasoning_content = chunk.replace("πspecial reasoning blockπ: ", "")
yield f"event: reasoning\ndata: {reasoning_content}\n\n"
elif chunk.startswith("πspecial clarification blockπ: "):
if isinstance(chunk, dict):
if "chunk_type" in chunk and chunk["chunk_type"] == "reasoning":
# Till the mobile team updates the app to support json format, we need to yield the chunk as a string
bwc_payload = {"reasoning": f"{chunk["header"]}\n\n{chunk["content"]}\n\n"}
# payload = {'header': chunk.get('header', ''), 'content': chunk.get('content', '')}
# yield f"event: reasoning\ndata: {json.dumps(payload)}\n\n"
yield f"event: reasoning\ndata: {json.dumps(bwc_payload)}\n\n"
else:
payload: Dict[str, Any] = {"chunk": chunk}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
elif isinstance(chunk, str) and chunk.startswith("πspecial clarification blockπ: "):
clarification_content = chunk.replace("πspecial clarification blockπ: ", "")
try:
clarification_data = json.loads(clarification_content)
@@ -155,19 +154,21 @@ async def get_answer(
log.warning(f"Failed to parse clarification JSON: {clarification_content}")
formatted_text = "I'm sorry, I can't understand your request in your workspace context. Can you be more specific?"
yield f"event: delta\ndata: {formatted_text}\n\n"
elif chunk.startswith("πspecial actions blockπ: "):
elif isinstance(chunk, str) and chunk.startswith("πspecial actions blockπ: "):
actions_content = chunk.replace("πspecial actions blockπ: ", "")
try:
actions_data = json.loads(actions_content)
yield f"event: actions\ndata: {json.dumps(actions_data)}\n\n"
except json.JSONDecodeError:
log.warning(f"Failed to parse actions JSON: {actions_content}")
yield f"event: delta\ndata: {chunk}\n\n"
payload = {"chunk": chunk}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
else:
yield f"event: delta\ndata: {chunk}\n\n"
payload = {"chunk": chunk}
yield f"event: delta\ndata: {json.dumps(payload)}\n\n"
# Prepare next iteration
next_chunk_task = asyncio.create_task(cast(Coroutine[None, None, str], base_iter.__anext__()))
next_chunk_task = asyncio.create_task(cast(Coroutine[None, None, Union[str, Dict[str, Any]]], base_iter.__anext__()))
finally:
heartbeat_stop.set()
if heartbeat_task and not heartbeat_task.done():
@@ -197,285 +198,6 @@ async def get_answer(
return JSONResponse(status_code=500, content={"detail": "Internal Server Error"})
@mobile_router.post("/initialize-chat/")
async def initialize_chat(
data: ChatInitializationRequestMobile,
db: AsyncSession = Depends(get_async_session),
token: HTTPAuthorizationCredentials = Depends(jwt_schema),
):
"""Initialize a new chat and return the chat_id immediately."""
try:
auth = await validate_jwt_token(token)
if not auth.user:
return JSONResponse(status_code=401, content={"detail": "Invalid User"})
user_id = auth.user.id
except Exception as e:
log.error(f"Error validating JWT: {e!s}")
return JSONResponse(status_code=401, content={"detail": "Invalid JWT"})
# Validate request data
validation_error = validate_chat_initialization(data)
if validation_error:
return JSONResponse(status_code=validation_error["status_code"], content={"detail": validation_error["detail"]})
# Initialize chat using standalone function (workspace details backfilled later in queue_answer)
result = await initialize_new_chat(
user_id=user_id,
db=db,
chat_id=data.chat_id,
is_project_chat=data.is_project_chat,
workspace_in_context=data.workspace_in_context,
workspace_id=data.workspace_id,
)
# Handle result from service layer
if result["success"]:
return JSONResponse(content={"chat_id": result["chat_id"]})
else:
# Map service error codes to HTTP status codes
status_code = 500 # default
if result.get("error_code") == "CHAT_EXISTS":
status_code = 409
elif result.get("error_code") == "CHAT_CREATION_FAILED":
status_code = 500
elif result.get("error_code") == "UNEXPECTED_ERROR":
status_code = 500
return JSONResponse(status_code=status_code, content={"detail": result["message"]})
@mobile_router.post("/queue-answer/")
async def queue_answer(
data: ChatRequestMobile,
db: AsyncSession = Depends(get_async_session),
token: HTTPAuthorizationCredentials = Depends(jwt_schema),
):
"""First phase of two-step streaming flow.
Persists the ChatRequest payload and returns a one-time stream token.
The token is simply the UUID of a freshly created *user* Message row so we
don't need a new table. The rest of the ChatRequest fields are stored in
a MessageFlowStep (tool_name="QUEUE", step_order=0) until the client
later redeems the token via /stream-answer/{token}."""
try:
auth = await validate_jwt_token(token)
if not auth.user:
return JSONResponse(status_code=401, content={"detail": "Invalid User"})
except Exception as e:
log.error(f"Error validating JWT: {e!s}")
return JSONResponse(status_code=401, content={"detail": "Invalid JWT"})
validation_error = validate_chat_request(data)
if validation_error:
return JSONResponse(status_code=validation_error["status_code"], content={"detail": validation_error["detail"]})
# 1. Create a new USER message that will serve as the token
if data.chat_id is None:
return JSONResponse(status_code=400, content={"detail": "chat_id is required. Call /initialize-chat/ first."})
## need to resolve ws id and slug from project id if it's project chat
if not data.workspace_id:
if data.is_project_chat and data.project_id:
log.info(f"Mobile Queue-answer: Resolving workspace_id from project_id: {data.project_id}")
resolved_workspace_id = await resolve_workspace_id_from_project_id(str(data.project_id))
# The DB may return an asyncpg UUID object. Convert safely to a standard uuid.UUID.
from uuid import UUID
workspace_id_to_use = UUID(str(resolved_workspace_id)) if resolved_workspace_id else None
log.info(f"Mobile Queue-answer: Resolved workspace_id: {workspace_id_to_use}")
else:
workspace_id_to_use = data.workspace_id
if not data.workspace_slug:
if workspace_id_to_use:
resolved_workspace_slug = await resolve_workspace_slug(workspace_id_to_use, data.workspace_slug)
else:
log.warning("Mobile Queue-answer: No workspace_id to resolve workspace_slug from")
resolved_workspace_slug = None
else:
resolved_workspace_slug = data.workspace_slug
from uuid import uuid4
token_id = uuid4()
user_message_res = await upsert_message(
message_id=token_id,
chat_id=data.chat_id, # type: ignore[arg-type]
content=data.query,
parsed_content=None,
user_type=UserTypeChoices.USER.value,
llm_model=data.llm,
workspace_slug=resolved_workspace_slug,
db=db,
)
if user_message_res.get("message") != "success":
return JSONResponse(status_code=500, content={"detail": "Failed to create message"})
# 2. Stash the full ChatRequest inside a flow-step row
from fastapi.encoders import jsonable_encoder
flow_step_payload = {
"step_order": 0,
"step_type": FlowStepType.TOOL.value,
"tool_name": "QUEUE",
"content": "queued chat request",
# Use FastAPI's encoder to turn UUIDs/datetimes into JSON-serialisable primitives
"execution_data": jsonable_encoder(data),
"is_planned": False, # QUEUE is not a planned action, it's an internal tool
"is_executed": False, # QUEUE is not executed by user
}
flow_res = await upsert_message_flow_steps(message_id=token_id, chat_id=data.chat_id, db=db, flow_steps=[flow_step_payload])
if flow_res.get("message") != "success":
return JSONResponse(status_code=500, content={"detail": "Failed to queue request"})
# Backfill chat record with workspace details (after successful message creation)
try:
from sqlalchemy import select
from pi.app.models.chat import Chat
# Get the current chat record to preserve existing data
chat_stmt = select(Chat).where(Chat.id == data.chat_id) # type: ignore[arg-type]
chat_result = await db.execute(chat_stmt)
existing_chat = chat_result.scalar_one_or_none()
if existing_chat:
# Update the chat with workspace details
if workspace_id_to_use is not None:
existing_chat.workspace_id = workspace_id_to_use
if resolved_workspace_slug is not None:
existing_chat.workspace_slug = resolved_workspace_slug
# Always update workspace_in_context as it's a required field in ChatRequest
existing_chat.workspace_in_context = data.workspace_in_context
await db.commit()
else:
log.warning(f"Chat {data.chat_id} not found for backfill")
except Exception as e:
log.error(f"Error backfilling chat workspace details: {e}")
# Don't fail the request if backfill fails
return {"stream_token": str(token_id)}
@mobile_router.get("/stream-answer/{token}")
async def stream_answer(
token: UUID4,
db: AsyncSession = Depends(get_async_session),
auth_token: HTTPAuthorizationCredentials = Depends(jwt_schema),
):
"""Second phase of two-step flow.
Looks up the queued ChatRequest by token (message_id), deletes the queue
entry, and then re-uses the existing /get-answer/ logic to start the SSE
stream via a pure GET endpoint.
Also handles REGENERATE when called with an existing user message ID
(no QUEUE flow step present)."""
try:
auth = await validate_jwt_token(auth_token)
if not auth.user:
return JSONResponse(status_code=401, content={"detail": "Invalid User"})
user_id = auth.user.id
except Exception:
return JSONResponse(status_code=401, content={"detail": "Invalid JWT"})
# Locate the queued flow step
from sqlalchemy import select
stmt = (
select(MessageFlowStep)
.where(MessageFlowStep.message_id == token) # type: ignore[arg-type]
.where(MessageFlowStep.tool_name == "QUEUE") # type: ignore[arg-type]
)
res = await db.execute(stmt)
flow_step: MessageFlowStep | None = res.scalar_one_or_none()
if not flow_step:
# REGENERATE FLOW: No QUEUE flow step means this is regenerate
log.info(f"Mobile: No QUEUE flow step found for token {token}. Treating as regenerate request.")
# 1. Get the user message
user_message = await get_message_by_id(db, token)
if not user_message:
log.warning(f"Mobile: Message not found for token {token}")
return JSONResponse(status_code=404, content={"detail": "Message not found"})
# 2. Mark old assistant response as replaced BEFORE generating new one
# This ensures retrieve_chat_history() won't include it in context for new generation
marked = await mark_assistant_response_as_replaced(db, user_message.id)
if marked:
log.info(f"Mobile: Marked old assistant response(s) as replaced for message {token}")
else:
log.info(f"Mobile: No existing assistant response found for message {token} (first generation or already replaced)")
# 4. Reconstruct ChatRequest from user message
try:
queued_request = await reconstruct_chat_request_from_message(db, user_message, user_id)
except Exception as e:
log.error(f"Mobile: Error reconstructing ChatRequest from message {token}: {e}")
return JSONResponse(status_code=500, content={"detail": "Failed to reconstruct request"})
# 5. Pass token_id so new assistant message reuses same user message as parent
try:
queued_request.context["token_id"] = str(token)
except Exception as e:
log.warning(f"Mobile: Failed to attach token_id to regenerate request context: {e!s}")
log.info(f"Mobile: Regenerating response for message {token}")
# 6. Stream new response (get_answer will call process_query_stream)
# When process_query_stream calls retrieve_chat_history,
# it will NOT see the old response because is_replaced=True
return await get_answer(data=queued_request, token=auth_token, db=db)
else:
# NORMAL FLOW: QUEUE flow step exists, this is first generation
# Parse the stored ChatRequest
try:
raw_data = flow_step.execution_data or {}
# Check if this is an OAuth message (has oauth_required field)
if flow_step.oauth_required:
# This is an OAuth message - check if OAuth is now complete
# by looking at the oauth_completed column
if flow_step.oauth_completed:
# OAuth is now complete! Process the request normally
log.info(f"Mobile: OAuth completed for message {token}. Processing request.")
# Control continues to the normal processing code below
else:
# OAuth is still required
return JSONResponse(
status_code=401,
content={
"detail": "OAuth authorization still required. Please complete OAuth authentication first.",
"error_code": "OAUTH_REQUIRED",
},
)
# Convert empty-string UUIDs to None so pydantic validation passes
for field in ["project_id", "workspace_id", "chat_id"]:
if field in raw_data and raw_data[field] == "":
raw_data[field] = None
raw_data["user_id"] = user_id
queued_request = ChatRequestMobile.parse_obj(raw_data)
except Exception as e:
log.error(f"Mobile: Malformed execution_data for token {token}: {e!s}")
return JSONResponse(status_code=500, content={"detail": "Corrupted queued request"})
# Consume the queue entry so the token is single-use (for normal flow)
await db.delete(flow_step)
await db.commit()
# Pass the token/message_id forward so downstream processing reuses the same Message row
try:
queued_request.context["token_id"] = str(token)
except Exception as e:
log.warning(f"Mobile: Failed to attach token_id to queued request context: {e!s}")
# Delegate to existing get_answer for streaming
return await get_answer(data=queued_request, token=auth_token, db=db)
@mobile_router.delete("/delete-chat/")
async def delete_chat(
data: DeleteChatRequestMobile,
@@ -617,8 +339,11 @@ async def get_chat_history(
"feedback": results["feedback"],
"reasoning": results.get("reasoning", ""),
"is_focus_enabled": results.get("is_focus_enabled", False),
"focus_entity_type": results.get("focus_entity_type", None),
"focus_entity_id": results.get("focus_entity_id", None),
"focus_project_id": results.get("focus_project_id", None),
"focus_workspace_id": results.get("focus_workspace_id", None),
"mode": results.get("mode", "ask"),
}
}
)
@@ -658,8 +383,11 @@ async def get_chat_history_object(
"feedback": results["feedback"],
"reasoning": results.get("reasoning", ""),
"is_focus_enabled": results.get("is_focus_enabled", False),
"focus_entity_type": results.get("focus_entity_type", None),
"focus_entity_id": results.get("focus_entity_id", None),
"focus_project_id": results.get("focus_project_id", None),
"focus_workspace_id": results.get("focus_workspace_id", None),
"mode": results.get("mode", "ask"),
}
}
)
apps/pi/pi/app/api/v1/helpers/batch_execution_helpers.py
-615
View File
@@ -1,615 +0,0 @@
"""
Helper functions for batch action execution in chat endpoints.
"""
import datetime
import logging
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Set
from uuid import UUID
from sqlalchemy import select
from sqlmodel.ext.asyncio.session import AsyncSession
from pi.app.models.enums import UserTypeChoices
from pi.app.models.message import Message
from pi.app.schemas.chat import ActionBatchExecutionRequest
from pi.app.schemas.chat import ArtifactData
from pi.services.chat.chat import PlaneChatBot
from pi.services.chat.helpers.batch_action_orchestrator import BatchActionOrchestrator
from pi.services.chat.helpers.batch_execution_context import BatchExecutionContext
from pi.services.chat.helpers.batch_execution_helpers import get_planned_actions_for_execution
from pi.services.chat.helpers.tool_utils import classify_tool
from pi.services.retrievers.pg_store.action_artifact import get_action_artifacts_by_ids
log = logging.getLogger(__name__)
# Registry to handle project/workspace id for entities during manual edits.
# HOW IT WORKS:
# 1. During planning, if project_id is missing, user gets clarification prompt
# 2. For workspace-scope entities, clarification includes "Workspace level" option
# 3. User selects either "Workspace level" or a specific project
# 4. Artifact stores project_id as either "__workspace_scope__" or a UUID
# 5. During execution, the tool checks project_id and routes accordingly:
# - "__workspace_scope__" → calls workspace-level API (e.g., create_workspace_page)
# - UUID → calls project-level API (e.g., create_project_page)
#
# EDITED ARTIFACTS:
# When user edits an artifact, the resolved project_id must be preserved.
# The _extract_project_id_from_artifact() helper extracts it from planning data.
WORKSPACE_SCOPE_ENTITIES: Set[str] = {
"page", # Pages: workspace wiki OR project pages
}
def _is_valid_uuid(uuid_string: str) -> bool:
"""Check if a string is a valid UUID format."""
try:
UUID(uuid_string)
return True
except (ValueError, TypeError):
return False
def _extract_project_id_from_artifact(artifact_data: Dict[str, Any]) -> Optional[str]:
"""Extract project_id from artifact data, checking multiple possible locations.
This handles the various ways project_id can be stored in artifact data:
1. In tool_args (most direct)
2. In planning_data.parameters (from planning phase)
3. In planning_data.parameters.project.id (nested structure)
Returns:
project_id as string (can be UUID or "__workspace_scope__"), or None if not found
"""
# Try tool_args first (most direct)
tool_args = artifact_data.get("tool_args", {})
if isinstance(tool_args, dict) and tool_args.get("project_id"):
return str(tool_args["project_id"])
# Try planning_data structure
planning_data = artifact_data.get("planning_data", {})
if isinstance(planning_data, dict):
parameters = planning_data.get("parameters", {})
if isinstance(parameters, dict):
# Direct project_id in parameters
if parameters.get("project_id"):
return str(parameters["project_id"])
# Nested project object with id
project_obj = parameters.get("project")
if isinstance(project_obj, dict) and project_obj.get("id"):
return str(project_obj["id"])
return None
def _extract_tool_args_from_artifact_data(artifact_data: dict, entity_type: str) -> dict:
"""
Extract and format tool arguments from artifact data for execution.
Handles different data formats:
1. Direct tool_args (from planning phase)
2. Nested parameters structure (from planning_data)
3. Flattened parameters (from UI edits)
"""
try:
# First, try direct tool_args (most reliable)
if "tool_args" in artifact_data:
tool_args = artifact_data["tool_args"]
if isinstance(tool_args, dict) and tool_args:
return tool_args
# Next, try planning_data.parameters structure
if "planning_data" in artifact_data:
planning_data = artifact_data["planning_data"]
if isinstance(planning_data, dict):
parameters = planning_data.get("parameters", {})
if isinstance(parameters, dict):
# Check for nested entity structure
if entity_type in parameters:
entity_params = parameters[entity_type]
if isinstance(entity_params, dict):
# Flatten entity parameters and merge with other parameters
flattened = dict(entity_params)
# Add other non-entity parameters
for key, value in parameters.items():
if key != entity_type:
flattened[key] = value
return flattened
else:
# Parameters are already flattened
return parameters
# Fallback: return empty dict if no valid args found
log.warning(f"No valid tool args found in artifact data for entity {entity_type}")
return {}
except Exception as e:
log.error(f"Error extracting tool args from artifact data: {e}")
return {}
async def validate_session_and_get_user(session: str) -> Optional[UUID]:
"""Validate session and return user ID if valid."""
try:
from pi.app.api.v1.dependencies import is_valid_session
auth = await is_valid_session(session)
if not auth.user:
return None
return auth.user.id
except Exception as e:
log.error(f"Error validating session: {e!s}")
return None
async def prepare_execution_data(request: ActionBatchExecutionRequest, user_id: UUID, db: AsyncSession) -> Optional[dict]:
"""Prepare execution data and return error response if validation fails."""
try:
# **ARTIFACT-BASED EXECUTION**
# Prefer explicit artifact_data from the client. If absent, fall back to executing
# all planned artifacts for this message to maintain backward compatibility.
artifact_items = request.artifact_data or []
if not artifact_items:
log.info("No artifact_data provided in request; attempting fallback to planned actions")
planned_for_message = await get_planned_actions_for_execution(request.message_id, request.chat_id, db)
if not planned_for_message:
log.error("No artifact_data provided and no planned actions found")
return None
# Build a default list of artifacts to execute (all planned, unedited)
tmp_items: list[ArtifactData] = []
for pa in planned_for_message:
aid = pa.get("artifact_id")
if not aid:
continue
try:
tmp_items.append(ArtifactData(artifact_id=UUID(aid), is_edited=False))
except Exception as e:
log.error(f"Invalid artifact_id in planned actions: {aid}: {e}")
if not tmp_items:
log.error("No valid artifact IDs found in planned actions for fallback execution")
return None
artifact_items = tmp_items
log.info(f"Processing artifacts: {len(artifact_items)} artifacts")
planned_actions: List[Dict[str, Any]] = []
for artifact_item in artifact_items:
log.info(f"Processing artifact: {artifact_item.artifact_id} (is_edited: {artifact_item.is_edited})")
# Get artifact metadata
artifacts = await get_action_artifacts_by_ids(db, [artifact_item.artifact_id])
if not artifacts:
log.error(f"Artifact {artifact_item.artifact_id} not found")
continue
artifact = artifacts[0]
entity_type = artifact.entity
entity_id = artifact.entity_id
original_action = artifact.action
if artifact_item.is_edited:
# **EDITED ARTIFACT** - Create ActionArtifactVersion and use its UUID
if not artifact_item.action_data:
log.error(f"Artifact {artifact_item.artifact_id} marked as edited but no action_data provided")
continue
# For edited artifacts, use the provided action_data directly
tool_args = artifact_item.action_data.copy()
current_project_id = tool_args.get("project_id")
# Check if project_id is missing or invalid (needs to be preserved from original)
needs_project_id = (
not current_project_id
or current_project_id == "NEEDS_CLARIFICATION"
or (isinstance(current_project_id, str) and current_project_id.startswith("<id of"))
)
if needs_project_id:
# Extract project_id from original artifact data
original_artifact_data = artifact.data or {}
original_project_id = _extract_project_id_from_artifact(original_artifact_data)
if original_project_id:
tool_args["project_id"] = original_project_id
else:
# No project_id found in original artifact - this might be okay for some entities
log.warning(f"No project_id found in original artifact {artifact_item.artifact_id} (entity: {entity_type})")
# Create ActionArtifactVersion and use its UUID as step_id
try:
from pi.services.actions.artifacts.utils import convert_uuids_to_strings
from pi.services.retrievers.pg_store.action_artifact import create_action_artifact_version
# Ensure all UUID objects are converted to strings for JSON serialization
serializable_tool_args = convert_uuids_to_strings(tool_args)
version = await create_action_artifact_version(
db=db,
artifact_id=artifact_item.artifact_id,
data=serializable_tool_args, # Use UUID-safe tool_args
change_type="manual_edit",
chat_id=request.chat_id,
message_id=request.message_id,
user_id=str(user_id),
)
if version:
step_id = str(version.id) # Use ActionArtifactVersion UUID as step_id
version_id = str(version.id)
log.info(f"Created ActionArtifactVersion {version.id} for edited artifact {artifact_item.artifact_id}")
else:
log.error(f"Failed to create ActionArtifactVersion for artifact {artifact_item.artifact_id}")
continue
except Exception as e:
log.error(f"Error creating ActionArtifactVersion for artifact {artifact_item.artifact_id}: {e}")
continue
log.info(f"Using edited artifact data for artifact {artifact_item.artifact_id}")
else:
# **NORMAL ARTIFACT** - Get the MessageFlowStep UUID
# Find the flow step for this artifact to get its UUID
planned_actions_for_artifact = await get_planned_actions_for_execution(request.message_id, request.chat_id, db)
# Find the planned action for this specific artifact
planned_action = None
for pa in planned_actions_for_artifact:
if pa.get("artifact_id") == str(artifact_item.artifact_id):
planned_action = pa
break
if not planned_action:
log.error(f"No planned action found for artifact {artifact_item.artifact_id}")
continue
# Use the flow step UUID as step_id
step_id_raw = planned_action.get("step_id")
if not step_id_raw or not _is_valid_uuid(str(step_id_raw)):
log.error(f"Invalid or missing step_id for artifact {artifact_item.artifact_id}: {step_id_raw}")
continue
step_id = str(step_id_raw)
# Get tool args from the planned action
tool_args = planned_action.get("args", {})
if not tool_args:
log.error(f"No tool args found for artifact {artifact_item.artifact_id}")
continue
version_id = None # Normal artifacts don't have versions
log.info(f"Using normal artifact data for artifact {artifact_item.artifact_id} with step_id {step_id}")
# Build tool name and args
tool_name = f"{entity_type}s_{original_action}"
# Validate and clean tool args
if not tool_args:
log.error(f"No tool args found for artifact {artifact_item.artifact_id}, entity: {entity_type}, action: {original_action}")
continue
args = {
**tool_args,
# Only add entity ID for update/delete actions
**(
{f"{entity_type}_id" if entity_type != "workitem" else "issue_id": str(entity_id)}
if original_action in ["update", "delete"] and entity_id
else {}
),
}
# Log the final args for debugging
log.info(f"Tool args for {tool_name}: {list(args.keys())}")
# Add to planned actions
planned_action_item = {
"step_id": step_id, # Now always a proper UUID
"step_order": len(planned_actions) + 1,
"tool_name": tool_name,
"args": args,
"action_summary": {"action": original_action, "entity": entity_type},
"tool_id": None,
"artifact_id": str(artifact_item.artifact_id),
"artifact_type": entity_type, # Add artifact type for response formatting
"sequence": len(planned_actions) + 1,
}
# Add version_id for edited artifacts
if artifact_item.is_edited and version_id:
planned_action_item["version_id"] = version_id
planned_action_item["is_edited"] = True
planned_actions.append(planned_action_item)
edit_type = "edited" if artifact_item.is_edited else "normal"
log.info(f"Added {tool_name} for artifact {artifact_item.artifact_id} ({edit_type})")
if not planned_actions:
log.error("No valid artifacts found")
return None
# Create appropriate query message based on artifact types
normal_count = sum(1 for action in planned_actions if not action.get("is_edited", False))
edited_count = sum(1 for action in planned_actions if action.get("is_edited", False))
if normal_count > 0 and edited_count > 0:
original_query = f"Mixed execution: {normal_count} normal + {edited_count} edited artifacts"
elif edited_count > 0:
original_query = f"Execute {edited_count} edited artifacts"
else:
original_query = f"Execute {normal_count} artifacts"
# Get OAuth token for the user
chatbot = PlaneChatBot()
access_token = request.access_token or await chatbot._get_oauth_token_for_user(db, str(user_id), str(request.workspace_id))
if not access_token:
log.error(f"No valid OAuth token found for user {user_id} and workspace {request.workspace_id}")
return None
# Get workspace context
from pi.app.api.v1.helpers.plane_sql_queries import get_workspace_slug
workspace_slug = await get_workspace_slug(str(request.workspace_id))
if not workspace_slug:
return None
# Get is_project_chat from chat record (with explicit commit/refresh to avoid race condition)
from sqlalchemy import select
from pi.app.models.chat import Chat
# Refresh the session to ensure we see any committed chat records
await db.commit()
stmt = select(Chat).where(Chat.id == request.chat_id) # type: ignore[arg-type]
result = await db.execute(stmt)
chat = result.scalar_one_or_none()
is_project_chat = chat.is_project_chat if chat else False
log.info(f"🔧 EXECUTION DATA PREP: chat_id={request.chat_id}, is_project_chat from DB (after commit)={is_project_chat}")
# Return successful execution data
return {
"planned_actions": planned_actions,
"original_query": original_query,
"access_token": access_token,
"workspace_slug": workspace_slug,
"user_id": user_id,
"workspace_id": request.workspace_id,
"message_id": request.message_id,
"chat_id": request.chat_id,
"is_project_chat": is_project_chat,
}
except Exception as e:
log.error(f"Error preparing execution data: {e}")
return None
async def execute_batch_actions(execution_data: dict, db: AsyncSession) -> BatchExecutionContext:
"""Execute batch actions using the orchestrator."""
try:
# Create execution context
context = BatchExecutionContext(
message_id=execution_data["message_id"],
chat_id=execution_data["chat_id"],
user_id=execution_data["user_id"],
workspace_id=execution_data["workspace_id"],
is_project_chat=execution_data.get("is_project_chat", False),
)
# Create orchestrator and execute batch
chatbot = PlaneChatBot("gpt-4.1", execution_data["access_token"])
orchestrator = BatchActionOrchestrator(chatbot, db)
context = await orchestrator.execute_planned_actions(
context,
execution_data["original_query"],
execution_data["planned_actions"],
execution_data["access_token"],
execution_data["workspace_slug"],
)
return context
except Exception as e:
log.error(f"Error executing batch actions: {e}")
# Create a failed context to return
context = BatchExecutionContext(
message_id=execution_data["message_id"],
chat_id=execution_data["chat_id"],
user_id=execution_data["user_id"],
workspace_id=execution_data["workspace_id"],
is_project_chat=execution_data.get("is_project_chat", False),
)
context.add_execution_failure("system", "execution_error", str(e))
return context
async def update_assistant_message_with_execution_results(message_id: UUID, chat_id: UUID, context: BatchExecutionContext, db: AsyncSession) -> None:
"""Update the assistant message with execution results."""
try:
# Get the assistant message that relates to this user message
# The assistant message should have the user message as its parent_id
filters = [Message.parent_id == message_id, Message.chat_id == chat_id, Message.user_type == UserTypeChoices.ASSISTANT.value]
stmt = select(Message).where(*filters) # type: ignore[arg-type]
result = await db.execute(stmt)
message = result.scalar_one_or_none()
if not message:
log.error(f"Assistant message for user message ID {message_id} not found")
return
# Frontend now gets execution details via structured data in chat history
# No longer appending execution results to assistant message content
except Exception as e:
log.error(f"Error updating assistant message with execution results: {e}")
if db:
await db.rollback()
def format_execution_response(context: BatchExecutionContext) -> Dict[str, Any]:
"""Format the execution response with clean, non-redundant structure."""
try:
# Base response structure
response: Dict[str, Any] = {
"action_summary": {
"total_planned": context.total_planned,
"completed": context.completed_count,
"failed": context.failed_count,
"duration_seconds": round((datetime.datetime.utcnow() - context.start_time).total_seconds(), 2),
},
}
# Add actions with consolidated entity information
if context.executed_actions:
response["actions"] = create_clean_actions_response(context.executed_actions)
return response
except Exception as e:
log.error(f"Error formatting execution response: {e}")
return {"action_summary": {"total_planned": 0, "completed": 0, "failed": 0, "duration_seconds": 0.0}, "actions": []}
def create_clean_actions_response(executed_actions: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Create clean action results without duplicate entity information.
Filters out retrieval/search tools and only includes actual modifying actions.
"""
clean_actions = []
for action in executed_actions:
# Extract action from tool name (e.g., "workitems_create" -> "create")
tool_name = action.get("tool_name", "")
# Filter out retrieval tools - only show actual actions to frontend
if tool_name:
is_retrieval, is_action = classify_tool(tool_name)
if is_retrieval and not is_action:
log.debug(f"Filtering out retrieval tool from actions response: {tool_name}")
continue
action_name = "unknown"
if tool_name and "_" in tool_name:
action_name = tool_name.split("_", 1)[1] # Get everything after the first underscore
elif tool_name:
action_name = tool_name
action_data = {
"action": action_name,
"artifact_type": action.get("artifact_type"),
"success": action.get("success"),
"executed_at": action.get("executed_at"),
"artifact_id": action.get("artifact_id"), # NEW: Include artifact ID
"sequence": action.get("sequence"), # NEW: Include planned step order
"version_number": action.get("version_number"), # NEW: Include version sequence number
}
if action.get("success"):
# For successful actions, include essential entity info
entity_info = action.get("entity_info")
if entity_info and isinstance(entity_info, dict):
# Only include the most important entity fields
essential_entity = {}
for field in ["entity_url", "entity_name", "entity_type", "entity_id"]:
if field in entity_info and entity_info[field]:
essential_entity[field] = entity_info[field]
# Include issue_identifier when available (for work-items)
if entity_info.get("issue_identifier"):
essential_entity["issue_identifier"] = entity_info["issue_identifier"]
if essential_entity:
action_data["entity"] = essential_entity
# Add project_identifier at the action root when derivable
project_identifier = None
try:
# Prefer deriving from identifiers present in entity_info
if entity_info.get("entity_type") == "project":
project_identifier = entity_info.get("entity_identifier") or entity_info.get("project_identifier")
elif entity_info.get("issue_identifier") and isinstance(entity_info.get("issue_identifier"), str):
ident = entity_info.get("issue_identifier")
if "-" in ident:
project_identifier = ident.split("-", 1)[0]
elif entity_info.get("entity_url") and isinstance(entity_info.get("entity_url"), str):
url = entity_info.get("entity_url")
# Attempt to parse /browse/PROJECT-SEQ/ pattern
if "/browse/" in url:
after = url.split("/browse/", 1)[1]
ident = after.split("/", 1)[0]
if "-" in ident:
project_identifier = ident.split("-", 1)[0]
except Exception:
project_identifier = None
if project_identifier:
action_data["project_identifier"] = project_identifier
# Extract the nice success message from the result
result = action.get("result", "")
if result:
# Extract the nice message that comes after "✅ " and before "\n\n"
nice_message = extract_success_message(result)
if nice_message:
action_data["message"] = nice_message
else:
# Fallback to generic messages
if "created" in result.lower():
action_data["message"] = "Created successfully"
elif "updated" in result.lower():
action_data["message"] = "Updated successfully"
else:
action_data["message"] = "Action completed successfully"
else:
# For failed actions, include error message
error = action.get("error", "")
if error:
action_data["error"] = error.split("\n")[0]
clean_actions.append(action_data)
return clean_actions
def extract_success_message(result: str) -> str:
"""Extract the nice success message from the tool result."""
if not result or not isinstance(result, str):
return ""
# Look for the pattern: "✅ [message]\n\n"
if "✅" in result:
lines = result.split("\n")
for line in lines:
line = line.strip()
if line.startswith("✅"):
# Remove the "✅ " prefix and return the message
message = line[2:].strip() # Remove "✅ " (2 characters)
return message
return ""
def get_status_message(context: BatchExecutionContext) -> str:
"""Generate a user-friendly status message."""
if context.status == "success":
return f"Successfully executed {context.completed_count} actions"
elif context.status == "partial":
return f"Partial success: {context.completed_count} actions completed, {context.failed_count} failed"
else:
return f"Batch execution failed: {context.failed_count} actions failed"
apps/pi/pi/app/api/v1/helpers/execution.py
+29
View File
@@ -0,0 +1,29 @@
from typing import Optional
from pydantic import UUID4
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.api.v1.helpers.message import _get_message
log = logger.getChild(__name__)
async def chosen_llm(db: AsyncSession, message_id: UUID4) -> Optional[str]:
"""
Retrieve the LLM model chosen for a specific message.
Args:
db: Database session
message_id: UUID of the message to retrieve the LLM model from
Returns:
The LLM model string if found, None otherwise
"""
message = await _get_message(db=db, message_id=message_id)
if message:
log.info(f"Chosen LLM model for message {message_id}: {message.llm_model}")
return message.llm_model
log.warning(f"No LLM model found for message {message_id}. Will use GPT-4.1 asdefault model.")
return None
apps/pi/pi/app/api/v1/helpers/plane_sql_queries.py
+49 -14
View File
@@ -593,6 +593,41 @@ async def search_module_by_name(name: str, project_id: Optional[str] = None, wor
return None
async def search_type_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Search for an issue type by name and return its details."""
query = """
SELECT it.id, it.name, pit.project_id
FROM issue_types it
JOIN project_issue_types pit ON it.id = pit.issue_type_id
JOIN projects p ON pit.project_id = p.id
JOIN workspaces w ON p.workspace_id = w.id
WHERE it.name ILIKE $1
AND it.deleted_at IS NULL
AND pit.deleted_at IS NULL
"""
params = [f"%{name}%"]
if project_id:
query += " AND pit.project_id = $2"
params.append(project_id)
if workspace_slug:
query += " AND w.slug = $3"
params.append(workspace_slug)
query += " LIMIT 1"
try:
result = await PlaneDBPool.fetchrow(query, tuple(params))
if result:
return {"id": str(result["id"]), "name": result["name"], "project_id": str(result["project_id"])}
return None
except Exception as e:
log.error(f"Error searching for issue type '{name}': {e}, query: {query}, project_id: {project_id}, workspace_slug: {workspace_slug}")
return None
async def search_project_by_name(
name: str,
workspace_slug: Optional[str] = None,
@@ -691,15 +726,15 @@ async def search_cycle_by_name(name: str, project_id: Optional[str] = None, work
return None
async def search_current_cycle(project_id: str, workspace_slug: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Search for the current active cycle in a project (where today's date falls within start_date and end_date).
async def search_current_cycle(project_id: str, workspace_slug: Optional[str] = None) -> List[Dict[str, Any]]:
"""Search for the current active cycles in a project (where today's date falls within start_date and end_date).
Args:
project_id: Project ID to search within (required)
workspace_slug: Workspace slug (optional filter)
Returns:
Dictionary with cycle details if found, None otherwise
List of dictionaries with cycle details, empty list if none found
"""
query = """
SELECT c.id, c.name, c.project_id, c.workspace_id, c.start_date, c.end_date, w.slug AS workspace_slug
@@ -709,8 +744,8 @@ async def search_current_cycle(project_id: str, workspace_slug: Optional[str] =
WHERE c.project_id = $1
AND c.deleted_at IS NULL
AND c.archived_at IS NULL
AND c.start_date <= CURRENT_DATE
AND c.end_date >= CURRENT_DATE
AND c.start_date::date <= CURRENT_DATE
AND c.end_date::date >= CURRENT_DATE
"""
params = [project_id]
@@ -721,12 +756,13 @@ async def search_current_cycle(project_id: str, workspace_slug: Optional[str] =
params.append(workspace_slug)
param_index += 1
query += " ORDER BY c.start_date DESC LIMIT 1"
query += " ORDER BY c.start_date DESC"
try:
result = await PlaneDBPool.fetchrow(query, tuple(params))
if result:
return {
results = await PlaneDBPool.fetch(query, tuple(params))
cycles = []
for result in results:
cycles.append({
"id": str(result["id"]),
"name": result["name"],
"project_id": str(result["project_id"]),
@@ -734,13 +770,13 @@ async def search_current_cycle(project_id: str, workspace_slug: Optional[str] =
"workspace_slug": result.get("workspace_slug"),
"start_date": str(result["start_date"]) if result["start_date"] else None,
"end_date": str(result["end_date"]) if result["end_date"] else None,
}
return None
})
return cycles
except Exception as e:
log.error(
f"Error searching for current cycle in project '{project_id}': {e}, query: {query}, project_id: {project_id}, workspace_slug: {workspace_slug}" # noqa E501
) # noqa: E501
return None
return []
async def search_label_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Optional[Dict[str, Any]]:
@@ -879,6 +915,7 @@ async def search_user_by_name(
WHERE ({where_name})
AND wm.deleted_at IS NULL
AND u.is_active IS TRUE
AND u.is_bot IS FALSE
"""
if workspace_slug:
@@ -887,10 +924,8 @@ async def search_user_by_name(
param_index += 1
query += " LIMIT 20"
log.info(f"Searching for user by name query: {query}, params: {params}")
try:
rows = await PlaneDBPool.fetch(query, tuple(params))
log.info(f"Searching for user by name rows: {rows}")
return [
{
"id": str(r["id"]),
apps/pi/pi/app/api/v2/endpoints/actions.py
+177 -49
View File
@@ -5,15 +5,19 @@ from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.api.v2.dependencies import cookie_schema
from pi.app.api.v2.helpers.batch_execution_helpers import execute_batch_actions
from pi.app.api.v2.helpers.batch_execution_helpers import format_execution_response
from pi.app.api.v2.helpers.batch_execution_helpers import prepare_execution_data
from pi.app.api.v2.helpers.batch_execution_helpers import update_assistant_message_with_execution_results
from pi.app.api.v2.helpers.batch_execution_helpers import validate_session_and_get_user
from pi.app.api.v2.dependencies import is_valid_session
# from pi.services.chat.helpers.action_execution_helpers import execute_batch_actions
# from pi.services.chat.helpers.action_execution_helpers import format_execution_response
# from pi.services.chat.helpers.action_execution_helpers import prepare_execution_data
# from pi.services.chat.helpers.action_execution_helpers import update_assistant_message_with_execution_results
# from pi.services.chat.helpers.action_execution_helpers import validate_session_and_get_user
from pi.app.schemas.chat import ActionBatchExecutionRequest
from pi.core.db.plane_pi.lifecycle import get_async_session
from pi.services.chat.helpers.batch_execution_helpers import get_original_user_query
from pi.services.chat.helpers.batch_execution_helpers import get_planned_actions_for_execution
from pi.services.chat.action_executor import BuildModeToolExecutor
from pi.services.chat.chat import PlaneChatBot
# from pi.services.chat.helpers.action_execution_helpers import get_original_user_query
log = logger.getChild("v2.actions")
router = APIRouter()
@@ -120,52 +124,31 @@ async def execute_batch_actions_endpoint(
"""
try:
# Validate session and get user
user_id = await validate_session_and_get_user(session)
if not user_id:
return JSONResponse(
status_code=401,
content={"detail": BATCH_EXECUTION_ERRORS["INVALID_SESSION"]},
)
auth = await is_valid_session(session)
if not auth.user:
return JSONResponse(status_code=401, content={"detail": "Invalid User"})
# Validate and prepare execution data
execution_data = await prepare_execution_data(request, user_id, db)
if not execution_data:
# Determine the specific error based on what failed
if not await get_planned_actions_for_execution(request.message_id, request.chat_id, db):
return JSONResponse(
status_code=404,
content={"detail": BATCH_EXECUTION_ERRORS["NO_PLANNED_ACTIONS"]},
)
elif not await get_original_user_query(request.message_id, db):
return JSONResponse(
status_code=404,
content={"detail": BATCH_EXECUTION_ERRORS["NO_ORIGINAL_QUERY"]},
)
else:
# OAuth or workspace issue
return JSONResponse(
status_code=401,
content={
"detail": BATCH_EXECUTION_ERRORS["OAUTH_REQUIRED"],
"error_code": "OAUTH_REQUIRED",
"workspace_id": str(request.workspace_id),
"user_id": str(user_id),
},
)
build_mode_tool_executor = BuildModeToolExecutor(chatbot=PlaneChatBot("gpt-4.1"), db=db)
result = await build_mode_tool_executor.execute(request, auth.user.id)
# Execute batch actions
context = await execute_batch_actions(execution_data, db)
# Check if service returned an error
if result.get("error"):
status_code = result.get("status_code", 500)
detail = result.get("detail", "Unknown error")
# Update the assistant message with execution results
await update_assistant_message_with_execution_results(
request.message_id,
request.chat_id,
context,
db,
)
# Build response content
content = {"detail": detail}
if "error_code" in result:
content["error_code"] = result["error_code"]
if "workspace_id" in result:
content["workspace_id"] = result["workspace_id"]
if "user_id" in result:
content["user_id"] = result["user_id"]
# Return appropriate response based on execution status
return JSONResponse(content=format_execution_response(context))
return JSONResponse(status_code=status_code, content=content)
# Return successful response
return JSONResponse(content=result)
except Exception as e:
log.error(f"Error in execute_batch_actions: {str(e)}")
@@ -173,3 +156,148 @@ async def execute_batch_actions_endpoint(
status_code=500,
content={"detail": BATCH_EXECUTION_ERRORS["INTERNAL_ERROR"]},
)
# async def execute_batch_actions_endpoint(
# request: ActionBatchExecutionRequest,
# db: AsyncSession = Depends(get_async_session),
# session: str = Depends(cookie_schema),
# ):
# """
# Execute all planned actions in a message as a batch using LLM orchestration.
# This endpoint takes planned actions from an AI response and executes them in batch
# using the Plane SDK. The execution results are stored and associated with the
# assistant message for context in future conversations.
# Execution Status Tracking:
# - When actions are planned: marked with is_executed=False in MessageFlowStep
# - When this endpoint is called: actions are marked with is_executed=True
# - Assistant messages are updated with execution results and entity information
# - Conversation history includes explicit text about executed vs. not executed actions
# This ensures complete context for follow-up questions and LLM understanding.
# Args:
# request: ActionBatchExecutionRequest containing:
# - workspace_id: UUID of workspace
# - chat_id: UUID of chat
# - message_id: UUID of message with planned actions
# - artifact_data: List of artifacts with action data
# - access_token: Optional OAuth token for API calls
# db: Database session (injected)
# session: Session cookie for authentication (injected)
# Returns:
# JSON response with:
# - execution_results: List of executed action results
# - errors: List of any errors encountered
# - summary: Execution summary statistics
# Status Codes:
# - 200: Actions executed successfully (even if some failed)
# - 401: Invalid authentication or OAuth required
# - 404: Message or planned actions not found
# - 500: Internal server error
# Example Request:
# POST /api/v2/actions/execute
# {
# "workspace_id": "abc-123",
# "chat_id": "chat-456",
# "message_id": "msg-789",
# "artifact_data": [
# {
# "artifact_id": "artifact-001",
# "is_edited": false,
# "action_data": {...}
# }
# ]
# }
# Example Response:
# {
# "execution_results": [
# {
# "action": "create_issue",
# "status": "success",
# "entity_id": "issue-123",
# "entity_name": "New bug report"
# }
# ],
# "errors": [],
# "summary": {
# "total": 1,
# "successful": 1,
# "failed": 0
# }
# }
# Notes:
# - Requires valid OAuth token for the workspace
# - Actions are executed in the order provided
# - Execution results are stored in the database
# - Failed actions don't stop execution of remaining actions
# - Deprecated V1 endpoint: POST /api/v1/chat/execute-action/
# Use Cases:
# - Execute issue creation/updates planned by AI
# - Execute project/cycle/module management actions
# - Execute bulk operations on multiple entities
# - Apply AI-suggested changes to workspace
# """
# try:
# # Validate session and get user
# user_id = await validate_session_and_get_user(session)
# if not user_id:
# return JSONResponse(
# status_code=401,
# content={"detail": BATCH_EXECUTION_ERRORS["INVALID_SESSION"]},
# )
# # Validate and prepare execution data
# execution_data = await prepare_execution_data(request, user_id, db)
# if not execution_data:
# # Determine the specific error based on what failed
# if not await get_planned_actions_for_execution(request.message_id, request.chat_id, db):
# return JSONResponse(
# status_code=404,
# content={"detail": BATCH_EXECUTION_ERRORS["NO_PLANNED_ACTIONS"]},
# )
# elif not await get_original_user_query(request.message_id, db):
# return JSONResponse(
# status_code=404,
# content={"detail": BATCH_EXECUTION_ERRORS["NO_ORIGINAL_QUERY"]},
# )
# else:
# # OAuth or workspace issue
# return JSONResponse(
# status_code=401,
# content={
# "detail": BATCH_EXECUTION_ERRORS["OAUTH_REQUIRED"],
# "error_code": "OAUTH_REQUIRED",
# "workspace_id": str(request.workspace_id),
# "user_id": str(user_id),
# },
# )
# # Execute batch actions
# context = await execute_batch_actions(execution_data, db)
# # Update the assistant message with execution results
# await update_assistant_message_with_execution_results(
# request.message_id,
# request.chat_id,
# context,
# db,
# )
# # Return appropriate response based on execution status
# return JSONResponse(content=format_execution_response(context))
# except Exception as e:
# log.error(f"Error in execute_batch_actions: {str(e)}")
# return JSONResponse(
# status_code=500,
# content={"detail": BATCH_EXECUTION_ERRORS["INTERNAL_ERROR"]},
# )
apps/pi/pi/app/api/v2/endpoints/responses.py
+25 -29
View File
@@ -22,9 +22,6 @@ from pi import logger
from pi.app.api.v1.endpoints._sse import normalize_error_chunk
from pi.app.api.v1.endpoints._sse import sse_done
from pi.app.api.v1.endpoints._sse import sse_event
from pi.app.api.v1.helpers.batch_execution_helpers import execute_batch_actions
from pi.app.api.v1.helpers.batch_execution_helpers import format_execution_response
from pi.app.api.v1.helpers.batch_execution_helpers import prepare_execution_data
from pi.app.api.v1.helpers.plane_sql_queries import resolve_workspace_id_from_project_id
from pi.app.api.v2.dependencies import cookie_schema
from pi.app.api.v2.dependencies import is_valid_session
@@ -40,9 +37,8 @@ from pi.app.utils.background_tasks import schedule_chat_search_upsert
from pi.app.utils.exceptions import SQLGenerationError
from pi.core.db.plane_pi.lifecycle import get_async_session
from pi.core.db.plane_pi.lifecycle import get_streaming_db_session
from pi.services.chat.action_executor import BuildModeToolExecutor
from pi.services.chat.chat import PlaneChatBot
from pi.services.chat.helpers.batch_execution_helpers import get_original_user_query
from pi.services.chat.helpers.batch_execution_helpers import get_planned_actions_for_execution
from pi.services.chat.helpers.tool_utils import format_clarification_as_text
from pi.services.chat.utils import resolve_workspace_slug
from pi.services.retrievers.pg_store.message import get_message_by_id
@@ -114,8 +110,11 @@ async def create_response_slack(data: ChatRequest, request: Request, db: AsyncSe
# Listen to all the stream chunks, join the chunks and return the complete response as JSON
async with get_streaming_db_session() as stream_db:
base_iter = chatbot.process_query_stream(data, db=stream_db)
base_iter = chatbot.process_chat_stream(data, db=stream_db)
async for chunk in base_iter:
if isinstance(chunk, dict):
# Currently only reasoning chunk is sent as dict.
continue
# Ignore all intermediate chunks
if chunk.startswith("πspecial reasoning block"):
continue
@@ -157,7 +156,10 @@ async def create_response_slack(data: ChatRequest, request: Request, db: AsyncSe
else:
return JSONResponse(status_code=400, content={"detail": "Invalid artifact_id format"})
execution_data = await prepare_execution_data(
# Execute batch actions using the service
service = BuildModeToolExecutor(chatbot=PlaneChatBot("gpt-4.1"), db=db)
result = await service.execute(
ActionBatchExecutionRequest(
workspace_id=workspace_id,
chat_id=chat_id,
@@ -166,29 +168,23 @@ async def create_response_slack(data: ChatRequest, request: Request, db: AsyncSe
access_token=access_token,
),
user_id,
db,
)
if not execution_data:
# Determine the specific error based on what failed
if not await get_planned_actions_for_execution(message_id, chat_id, db):
return JSONResponse(status_code=404, content={"detail": BATCH_EXECUTION_ERRORS["NO_PLANNED_ACTIONS"]})
elif not await get_original_user_query(message_id, db):
return JSONResponse(status_code=404, content={"detail": BATCH_EXECUTION_ERRORS["NO_ORIGINAL_QUERY"]})
else:
# OAuth or workspace issue
return JSONResponse(
status_code=401,
content={
"detail": BATCH_EXECUTION_ERRORS["OAUTH_REQUIRED"],
"error_code": "OAUTH_REQUIRED",
"workspace_id": str(data.workspace_id),
"user_id": str(data.user_id),
},
)
# Execute batch actions
context = await execute_batch_actions(execution_data, db)
formatted_context = format_execution_response(context)
# Check if service returned an error
if result.get("error"):
status_code = result.get("status_code", 500)
detail = result.get("detail", "Unknown error")
content = {"detail": detail}
if "error_code" in result:
content["error_code"] = result["error_code"]
if "workspace_id" in result:
content["workspace_id"] = result["workspace_id"]
if "user_id" in result:
content["user_id"] = result["user_id"]
return JSONResponse(status_code=status_code, content=content)
# Extract the formatted context from successful result
formatted_context = result
return JSONResponse(
status_code=200,
@@ -267,7 +263,7 @@ async def create_response_stream(data: ChatRequest, session: str = Depends(cooki
with contextlib.suppress(asyncio.CancelledError):
await asyncio.sleep(10)
base_iter = chatbot.process_query_stream(data, db=stream_db)
base_iter = chatbot.process_chat_stream(data, db=stream_db)
next_chunk_task: asyncio.Task[str] = asyncio.create_task(cast(Coroutine[None, None, str], base_iter.__anext__()))
# Start initial heartbeat timer
apps/pi/pi/app/api/v2/helpers/batch_execution_helpers.py
-624
View File
@@ -1,624 +0,0 @@
"""
Helper functions for batch action execution in chat endpoints.
"""
import datetime
import logging
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Set
from uuid import UUID
from sqlalchemy import select
from sqlmodel.ext.asyncio.session import AsyncSession
from pi.app.models.enums import UserTypeChoices
from pi.app.models.message import Message
from pi.app.schemas.chat import ActionBatchExecutionRequest
from pi.app.schemas.chat import ArtifactData
from pi.services.chat.chat import PlaneChatBot
from pi.services.chat.helpers.batch_action_orchestrator import BatchActionOrchestrator
from pi.services.chat.helpers.batch_execution_context import BatchExecutionContext
from pi.services.chat.helpers.batch_execution_helpers import get_planned_actions_for_execution
from pi.services.chat.helpers.tool_utils import classify_tool
from pi.services.retrievers.pg_store.action_artifact import get_action_artifacts_by_ids
log = logging.getLogger(__name__)
# Registry to handle project/workspace id for entities during manual edits.
# HOW IT WORKS:
# 1. During planning, if project_id is missing, user gets clarification prompt
# 2. For workspace-scope entities, clarification includes "Workspace level" option
# 3. User selects either "Workspace level" or a specific project
# 4. Artifact stores project_id as either "__workspace_scope__" or a UUID
# 5. During execution, the tool checks project_id and routes accordingly:
# - "__workspace_scope__" → calls workspace-level API (e.g., create_workspace_page)
# - UUID → calls project-level API (e.g., create_project_page)
#
# EDITED ARTIFACTS:
# When user edits an artifact, the resolved project_id must be preserved.
# The _extract_project_id_from_artifact() helper extracts it from planning data.
WORKSPACE_SCOPE_ENTITIES: Set[str] = {
"page", # Pages: workspace wiki OR project pages
}
def _is_valid_uuid(uuid_string: str) -> bool:
"""Check if a string is a valid UUID format."""
try:
UUID(uuid_string)
return True
except (ValueError, TypeError):
return False
def _extract_project_id_from_artifact(artifact_data: Dict[str, Any]) -> Optional[str]:
"""Extract project_id from artifact data, checking multiple possible locations.
This handles the various ways project_id can be stored in artifact data:
1. In tool_args (most direct)
2. In planning_data.parameters (from planning phase)
3. In planning_data.parameters.project.id (nested structure)
Returns:
project_id as string (can be UUID or "__workspace_scope__"), or None if not found
"""
# Try tool_args first (most direct)
tool_args = artifact_data.get("tool_args", {})
if isinstance(tool_args, dict) and tool_args.get("project_id"):
return str(tool_args["project_id"])
# Try planning_data structure
planning_data = artifact_data.get("planning_data", {})
if isinstance(planning_data, dict):
parameters = planning_data.get("parameters", {})
if isinstance(parameters, dict):
# Direct project_id in parameters
if parameters.get("project_id"):
return str(parameters["project_id"])
# Nested project object with id
project_obj = parameters.get("project")
if isinstance(project_obj, dict) and project_obj.get("id"):
return str(project_obj["id"])
return None
def _extract_tool_args_from_artifact_data(artifact_data: dict, entity_type: str) -> dict:
"""
Extract and format tool arguments from artifact data for execution.
Handles different data formats:
1. Direct tool_args (from planning phase)
2. Nested parameters structure (from planning_data)
3. Flattened parameters (from UI edits)
"""
try:
# First, try direct tool_args (most reliable)
if "tool_args" in artifact_data:
tool_args = artifact_data["tool_args"]
if isinstance(tool_args, dict) and tool_args:
return tool_args
# Next, try planning_data.parameters structure
if "planning_data" in artifact_data:
planning_data = artifact_data["planning_data"]
if isinstance(planning_data, dict):
parameters = planning_data.get("parameters", {})
if isinstance(parameters, dict):
# Check for nested entity structure
if entity_type in parameters:
entity_params = parameters[entity_type]
if isinstance(entity_params, dict):
# Flatten entity parameters and merge with other parameters
flattened = dict(entity_params)
# Add other non-entity parameters
for key, value in parameters.items():
if key != entity_type:
flattened[key] = value
return flattened
else:
# Parameters are already flattened
return parameters
# Fallback: return empty dict if no valid args found
log.warning(f"No valid tool args found in artifact data for entity {entity_type}")
return {}
except Exception as e:
log.error(f"Error extracting tool args from artifact data: {e}")
return {}
async def validate_session_and_get_user(session: str) -> Optional[UUID]:
"""Validate session and return user ID if valid."""
try:
from pi.app.api.v2.dependencies import is_valid_session
auth = await is_valid_session(session)
if not auth.user:
return None
return auth.user.id
except Exception as e:
log.error(f"Error validating session: {e!s}")
return None
async def prepare_execution_data(request: ActionBatchExecutionRequest, user_id: UUID, db: AsyncSession) -> Optional[dict]:
"""Prepare execution data and return error response if validation fails."""
try:
# **ARTIFACT-BASED EXECUTION**
# Prefer explicit artifact_data from the client. If absent, fall back to executing
# all planned artifacts for this message to maintain backward compatibility.
artifact_items = request.artifact_data or []
if not artifact_items:
log.info("No artifact_data provided in request; attempting fallback to planned actions")
planned_for_message = await get_planned_actions_for_execution(request.message_id, request.chat_id, db)
if not planned_for_message:
log.error("No artifact_data provided and no planned actions found")
return None
# Build a default list of artifacts to execute (all planned, unedited)
tmp_items: list[ArtifactData] = []
for pa in planned_for_message:
aid = pa.get("artifact_id")
if not aid:
continue
try:
tmp_items.append(ArtifactData(artifact_id=UUID(aid), is_edited=False))
except Exception as e:
log.error(f"Invalid artifact_id in planned actions: {aid}: {e}")
if not tmp_items:
log.error("No valid artifact IDs found in planned actions for fallback execution")
return None
artifact_items = tmp_items
log.info(f"Processing artifacts: {len(artifact_items)} artifacts")
planned_actions: List[Dict[str, Any]] = []
for artifact_item in artifact_items:
log.info(f"Processing artifact: {artifact_item.artifact_id} (is_edited: {artifact_item.is_edited})")
# Get artifact metadata
artifacts = await get_action_artifacts_by_ids(db, [artifact_item.artifact_id])
if not artifacts:
log.error(f"Artifact {artifact_item.artifact_id} not found")
continue
artifact = artifacts[0]
entity_type = artifact.entity
entity_id = artifact.entity_id
original_action = artifact.action
if artifact_item.is_edited:
# **EDITED ARTIFACT** - Create ActionArtifactVersion and use its UUID
if not artifact_item.action_data:
log.error(f"Artifact {artifact_item.artifact_id} marked as edited but no action_data provided")
continue
# For edited artifacts, use the provided action_data directly
tool_args = artifact_item.action_data.copy()
current_project_id = tool_args.get("project_id")
# Check if project_id is missing or invalid (needs to be preserved from original)
needs_project_id = (
not current_project_id
or current_project_id == "NEEDS_CLARIFICATION"
or (isinstance(current_project_id, str) and current_project_id.startswith("<id of"))
)
if needs_project_id:
# Extract project_id from original artifact data
original_artifact_data = artifact.data or {}
original_project_id = _extract_project_id_from_artifact(original_artifact_data)
if original_project_id:
tool_args["project_id"] = original_project_id
else:
# No project_id found in original artifact - this might be okay for some entities
log.warning(f"No project_id found in original artifact {artifact_item.artifact_id} (entity: {entity_type})")
# Resolve project_id to full project object for UI display
try:
from pi.services.actions.artifacts.utils import resolve_project_id_to_object
tool_args = await resolve_project_id_to_object(tool_args)
log.info(f"Resolved project data for artifact {artifact_item.artifact_id}")
except Exception as e:
log.error(f"Error resolving project_id for artifact {artifact_item.artifact_id}: {e}")
# Create ActionArtifactVersion and use its UUID as step_id
try:
from pi.services.actions.artifacts.utils import convert_uuids_to_strings
from pi.services.retrievers.pg_store.action_artifact import create_action_artifact_version
# Ensure all UUID objects are converted to strings for JSON serialization
serializable_tool_args = convert_uuids_to_strings(tool_args)
version = await create_action_artifact_version(
db=db,
artifact_id=artifact_item.artifact_id,
data=serializable_tool_args, # Use UUID-safe tool_args
change_type="manual_edit",
chat_id=request.chat_id,
message_id=request.message_id,
user_id=str(user_id),
)
if version:
step_id = str(version.id) # Use ActionArtifactVersion UUID as step_id
version_id = str(version.id)
log.info(f"Created ActionArtifactVersion {version.id} for edited artifact {artifact_item.artifact_id}")
else:
log.error(f"Failed to create ActionArtifactVersion for artifact {artifact_item.artifact_id}")
continue
except Exception as e:
log.error(f"Error creating ActionArtifactVersion for artifact {artifact_item.artifact_id}: {e}")
continue
log.info(f"Using edited artifact data for artifact {artifact_item.artifact_id}")
else:
# **NORMAL ARTIFACT** - Get the MessageFlowStep UUID
# Find the flow step for this artifact to get its UUID
planned_actions_for_artifact = await get_planned_actions_for_execution(request.message_id, request.chat_id, db)
# Find the planned action for this specific artifact
planned_action = None
for pa in planned_actions_for_artifact:
if pa.get("artifact_id") == str(artifact_item.artifact_id):
planned_action = pa
break
if not planned_action:
log.error(f"No planned action found for artifact {artifact_item.artifact_id}")
continue
# Use the flow step UUID as step_id
step_id_raw = planned_action.get("step_id")
if not step_id_raw or not _is_valid_uuid(str(step_id_raw)):
log.error(f"Invalid or missing step_id for artifact {artifact_item.artifact_id}: {step_id_raw}")
continue
step_id = str(step_id_raw)
# Get tool args from the planned action
tool_args = planned_action.get("args", {})
if not tool_args:
log.error(f"No tool args found for artifact {artifact_item.artifact_id}")
continue
version_id = None # Normal artifacts don't have versions
log.info(f"Using normal artifact data for artifact {artifact_item.artifact_id} with step_id {step_id}")
# Build tool name and args
tool_name = f"{entity_type}s_{original_action}"
# Validate and clean tool args
if not tool_args:
log.error(f"No tool args found for artifact {artifact_item.artifact_id}, entity: {entity_type}, action: {original_action}")
continue
args = {
**tool_args,
# Only add entity ID for update/delete actions
**(
{f"{entity_type}_id" if entity_type != "workitem" else "issue_id": str(entity_id)}
if original_action in ["update", "delete"] and entity_id
else {}
),
}
# Log the final args for debugging
log.info(f"Tool args for {tool_name}: {list(args.keys())}")
# Add to planned actions
planned_action_item = {
"step_id": step_id, # Now always a proper UUID
"step_order": len(planned_actions) + 1,
"tool_name": tool_name,
"args": args,
"action_summary": {"action": original_action, "entity": entity_type},
"tool_id": None,
"artifact_id": str(artifact_item.artifact_id),
"artifact_type": entity_type, # Add artifact type for response formatting
"sequence": len(planned_actions) + 1,
}
# Add version_id for edited artifacts
if artifact_item.is_edited and version_id:
planned_action_item["version_id"] = version_id
planned_action_item["is_edited"] = True
planned_actions.append(planned_action_item)
edit_type = "edited" if artifact_item.is_edited else "normal"
log.info(f"Added {tool_name} for artifact {artifact_item.artifact_id} ({edit_type})")
if not planned_actions:
log.error("No valid artifacts found")
return None
# Create appropriate query message based on artifact types
normal_count = sum(1 for action in planned_actions if not action.get("is_edited", False))
edited_count = sum(1 for action in planned_actions if action.get("is_edited", False))
if normal_count > 0 and edited_count > 0:
original_query = f"Mixed execution: {normal_count} normal + {edited_count} edited artifacts"
elif edited_count > 0:
original_query = f"Execute {edited_count} edited artifacts"
else:
original_query = f"Execute {normal_count} artifacts"
# Get OAuth token for the user
chatbot = PlaneChatBot()
access_token = request.access_token or await chatbot._get_oauth_token_for_user(db, str(user_id), str(request.workspace_id))
if not access_token:
log.error(f"No valid OAuth token found for user {user_id} and workspace {request.workspace_id}")
return None
# Get workspace context
from pi.app.api.v1.helpers.plane_sql_queries import get_workspace_slug
workspace_slug = await get_workspace_slug(str(request.workspace_id))
if not workspace_slug:
return None
# Get is_project_chat from chat record (with explicit commit/refresh to avoid race condition)
from sqlalchemy import select
from pi.app.models.chat import Chat
# Refresh the session to ensure we see any committed chat records
await db.commit()
stmt = select(Chat).where(Chat.id == request.chat_id) # type: ignore[arg-type]
result = await db.execute(stmt)
chat = result.scalar_one_or_none()
is_project_chat = chat.is_project_chat if chat else False
log.info(f"🔧 EXECUTION DATA PREP: chat_id={request.chat_id}, is_project_chat from DB (after commit)={is_project_chat}")
# Return successful execution data
return {
"planned_actions": planned_actions,
"original_query": original_query,
"access_token": access_token,
"workspace_slug": workspace_slug,
"user_id": user_id,
"workspace_id": request.workspace_id,
"message_id": request.message_id,
"chat_id": request.chat_id,
"is_project_chat": is_project_chat,
}
except Exception as e:
log.error(f"Error preparing execution data: {e}")
return None
async def execute_batch_actions(execution_data: dict, db: AsyncSession) -> BatchExecutionContext:
"""Execute batch actions using the orchestrator."""
try:
# Create execution context
context = BatchExecutionContext(
message_id=execution_data["message_id"],
chat_id=execution_data["chat_id"],
user_id=execution_data["user_id"],
workspace_id=execution_data["workspace_id"],
is_project_chat=execution_data.get("is_project_chat", False),
)
# Create orchestrator and execute batch
chatbot = PlaneChatBot("gpt-4.1", execution_data["access_token"])
orchestrator = BatchActionOrchestrator(chatbot, db)
context = await orchestrator.execute_planned_actions(
context,
execution_data["original_query"],
execution_data["planned_actions"],
execution_data["access_token"],
execution_data["workspace_slug"],
)
return context
except Exception as e:
log.error(f"Error executing batch actions: {e}")
# Create a failed context to return
context = BatchExecutionContext(
message_id=execution_data["message_id"],
chat_id=execution_data["chat_id"],
user_id=execution_data["user_id"],
workspace_id=execution_data["workspace_id"],
is_project_chat=execution_data.get("is_project_chat", False),
)
context.add_execution_failure("system", "execution_error", str(e))
return context
async def update_assistant_message_with_execution_results(message_id: UUID, chat_id: UUID, context: BatchExecutionContext, db: AsyncSession) -> None:
"""Update the assistant message with execution results."""
try:
# Get the assistant message that relates to this user message
# The assistant message should have the user message as its parent_id
filters = [Message.parent_id == message_id, Message.chat_id == chat_id, Message.user_type == UserTypeChoices.ASSISTANT.value]
stmt = select(Message).where(*filters) # type: ignore[arg-type]
result = await db.execute(stmt)
message = result.scalar_one_or_none()
if not message:
log.error(f"Assistant message for user message ID {message_id} not found")
return
# Frontend now gets execution details via structured data in chat history
# No longer appending execution results to assistant message content
except Exception as e:
log.error(f"Error updating assistant message with execution results: {e}")
if db:
await db.rollback()
def format_execution_response(context: BatchExecutionContext) -> Dict[str, Any]:
"""Format the execution response with clean, non-redundant structure."""
try:
# Base response structure
response: Dict[str, Any] = {
"action_summary": {
"total_planned": context.total_planned,
"completed": context.completed_count,
"failed": context.failed_count,
"duration_seconds": round((datetime.datetime.utcnow() - context.start_time).total_seconds(), 2),
},
}
# Add actions with consolidated entity information
if context.executed_actions:
response["actions"] = create_clean_actions_response(context.executed_actions)
return response
except Exception as e:
log.error(f"Error formatting execution response: {e}")
return {"action_summary": {"total_planned": 0, "completed": 0, "failed": 0, "duration_seconds": 0.0}, "actions": []}
def create_clean_actions_response(executed_actions: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Create clean action results without duplicate entity information.
Filters out retrieval/search tools and only includes actual modifying actions.
"""
clean_actions = []
for action in executed_actions:
# Extract action from tool name (e.g., "workitems_create" -> "create")
tool_name = action.get("tool_name", "")
# Filter out retrieval tools - only show actual actions to frontend
if tool_name:
is_retrieval, is_action = classify_tool(tool_name)
if is_retrieval and not is_action:
log.debug(f"Filtering out retrieval tool from actions response: {tool_name}")
continue
action_name = "unknown"
if tool_name and "_" in tool_name:
action_name = tool_name.split("_", 1)[1] # Get everything after the first underscore
elif tool_name:
action_name = tool_name
action_data = {
"action": action_name,
"artifact_type": action.get("artifact_type"),
"success": action.get("success"),
"executed_at": action.get("executed_at"),
"artifact_id": action.get("artifact_id"), # NEW: Include artifact ID
"sequence": action.get("sequence"), # NEW: Include planned step order
"version_number": action.get("version_number"), # NEW: Include version sequence number
}
if action.get("success"):
# For successful actions, include essential entity info
entity_info = action.get("entity_info")
if entity_info and isinstance(entity_info, dict):
# Only include the most important entity fields
essential_entity = {}
for field in ["entity_url", "entity_name", "entity_type", "entity_id"]:
if field in entity_info and entity_info[field]:
essential_entity[field] = entity_info[field]
# Include issue_identifier when available (for work-items)
if entity_info.get("issue_identifier"):
essential_entity["issue_identifier"] = entity_info["issue_identifier"]
if essential_entity:
action_data["entity"] = essential_entity
# Add project_identifier at the action root when derivable
project_identifier = None
try:
# Prefer deriving from identifiers present in entity_info
if entity_info.get("entity_type") == "project":
project_identifier = entity_info.get("entity_identifier") or entity_info.get("project_identifier")
elif entity_info.get("issue_identifier") and isinstance(entity_info.get("issue_identifier"), str):
ident = entity_info.get("issue_identifier")
if "-" in ident:
project_identifier = ident.split("-", 1)[0]
elif entity_info.get("entity_url") and isinstance(entity_info.get("entity_url"), str):
url = entity_info.get("entity_url")
# Attempt to parse /browse/PROJECT-SEQ/ pattern
if "/browse/" in url:
after = url.split("/browse/", 1)[1]
ident = after.split("/", 1)[0]
if "-" in ident:
project_identifier = ident.split("-", 1)[0]
except Exception:
project_identifier = None
if project_identifier:
action_data["project_identifier"] = project_identifier
# Extract the nice success message from the result
result = action.get("result", "")
if result:
# Extract the nice message that comes after "✅ " and before "\n\n"
nice_message = extract_success_message(result)
if nice_message:
action_data["message"] = nice_message
else:
# Fallback to generic messages
if "created" in result.lower():
action_data["message"] = "Created successfully"
elif "updated" in result.lower():
action_data["message"] = "Updated successfully"
else:
action_data["message"] = "Action completed successfully"
else:
# For failed actions, include error message
error = action.get("error", "")
if error:
action_data["error"] = error[:200] + "..." if len(error) > 200 else error
clean_actions.append(action_data)
return clean_actions
def extract_success_message(result: str) -> str:
"""Extract the nice success message from the tool result."""
if not result or not isinstance(result, str):
return ""
# Look for the pattern: "✅ [message]\n\n"
if "✅" in result:
lines = result.split("\n")
for line in lines:
line = line.strip()
if line.startswith("✅"):
# Remove the "✅ " prefix and return the message
message = line[2:].strip() # Remove "✅ " (2 characters)
return message
return ""
def get_status_message(context: BatchExecutionContext) -> str:
"""Generate a user-friendly status message."""
if context.status == "success":
return f"Successfully executed {context.completed_count} actions"
elif context.status == "partial":
return f"Partial success: {context.completed_count} actions completed, {context.failed_count} failed"
else:
return f"Batch execution failed: {context.failed_count} actions failed"
apps/pi/pi/app/models/chat.py
+15
View File
@@ -12,6 +12,7 @@ from sqlalchemy.dialects.postgresql import UUID
from sqlmodel import Field
from pi.app.models.base import BaseModel
from pi.app.models.enums import FocusEntityType
class Chat(BaseModel, table=True):
@@ -34,9 +35,23 @@ class UserChatPreference(BaseModel, table=True):
# Fields
is_focus_enabled: bool = Field(default=True, nullable=False)
# Polymorphic focus context - replaces focus_project_id and focus_workspace_id
focus_entity_type: Optional[str] = Field(default=None, nullable=True, max_length=50)
focus_entity_id: Optional[uuid.UUID] = Field(default=None, nullable=True)
# Legacy fields - kept for backward compatibility during migration
focus_project_id: Optional[uuid.UUID] = Field(default=None, nullable=True)
focus_workspace_id: Optional[uuid.UUID] = Field(default=None, nullable=True)
mode: str = Field(default="ask", nullable=False, max_length=10) # "ask" or "build"
user_id: uuid.UUID = Field(default=None, nullable=False)
# Foreign keys
chat_id: uuid.UUID = Field(sa_column=Column(UUID(as_uuid=True), ForeignKey("chats.id", name="fk_user_chat_preferences_chat_id"), nullable=False))
def get_focus_entity_type_enum(self) -> Optional[FocusEntityType]:
"""Get focus_entity_type as enum if valid, otherwise None."""
if self.focus_entity_type:
try:
return FocusEntityType(self.focus_entity_type)
except ValueError:
return None
return None
apps/pi/pi/app/models/enums.py
+19
View File
@@ -76,3 +76,22 @@ class ExecutionStatus(str, Enum):
def __str__(self):
return self.value
class FocusEntityType(str, Enum):
"""Entity types that can be used as focus context in chat preferences."""
WORKSPACE = "workspace"
PROJECT = "project"
CYCLE = "cycle"
MODULE = "module"
INITIATIVE = "initiative"
TEAMSPACE = "teamspace"
PAGE = "page"
PROJECT_PAGE = "project_page"
WIKI = "wiki"
EPIC = "epic"
WORKITEM = "workitem"
def __str__(self):
return self.value
apps/pi/pi/app/models/message.py
+3
View File
@@ -33,6 +33,9 @@ class Message(BaseModel, table=True):
reasoning: Optional[str] = Field(default=None, nullable=True)
user_type: str = Field(sa_column=Column(String(50), nullable=False, default=UserTypeChoices.USER.value))
workspace_slug: Optional[str] = Field(default=None, nullable=True, max_length=255)
source: Optional[str] = Field(
default=None, nullable=True, max_length=50, description="Source of the chat request (e.g., 'web', 'mobile', 'api', 'agent)"
)
# Foreign keys
parent_id: Optional[uuid.UUID] = Field(
apps/pi/pi/app/schemas/chat.py
+10
View File
@@ -3,6 +3,7 @@ from enum import Enum
from typing import Any
from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from pydantic import UUID4
@@ -52,9 +53,14 @@ class ChatRequest(BaseModel):
is_temp: bool
workspace_in_context: bool
# is_reasoning: bool = False
# New polymorphic focus context fields
focus_entity_type: Optional[str] = Field(default=None, description="Type of focus entity (workspace, project, cycle, module, etc.)")
focus_entity_id: Optional[UUID4] = Field(default=None, description="ID of the focus entity")
# Legacy fields (for backward compatibility)
workspace_id: UUID4 | None = "" # type: ignore
project_id: UUID4 | None = "" # type: ignore
context: dict[str, Any]
mode: Literal["ask", "build"] = "ask" # ask: normal chat, build: planning mode
is_project_chat: Optional[bool] = False
pi_sidebar_open: Optional[bool] = False
workspace_slug: Optional[str] = None
@@ -69,6 +75,10 @@ class ChatInitializationRequest(BaseModel):
user_id: Optional[UUID4] = None
chat_id: Optional[UUID4] = None
workspace_in_context: bool = False
# New polymorphic focus context fields
focus_entity_type: Optional[str] = Field(default=None, description="Type of focus entity (workspace, project, cycle, module, etc.)")
focus_entity_id: Optional[UUID4] = Field(default=None, description="ID of the focus entity")
# Legacy fields (for backward compatibility)
workspace_id: UUID4 | None = None
project_id: UUID4 | None = None
is_project_chat: Optional[bool] = False
apps/pi/pi/config.py
+8 -6
View File
@@ -8,7 +8,6 @@ from typing import Optional
import colorlog
from dotenv import load_dotenv
from pythonjsonlogger.json import JsonFormatter
load_dotenv()
@@ -197,8 +196,8 @@ class LLMConfig:
)
CONTEXT_OFF_TEMPERATURE: float = 0.6
OPENAI_RANDOM_SEED: int = 314
LITE_LLM_HOST: str = field(default_factory=lambda: os.getenv("LITE_LLM_HOST", "https://litellm.plane.town"))
LITE_LLM_API_KEY: str = field(default_factory=lambda: os.getenv("LITE_LLM_API_KEY", "sk-uFlW27K1uuF9hekh13QT9A"))
LITE_LLM_HOST: str = field(default_factory=lambda: os.getenv("LITE_LLM_HOST", ""))
LITE_LLM_API_KEY: str = field(default_factory=lambda: os.getenv("LITE_LLM_API_KEY", ""))
ENABLE_MODEL_VERIFICATION_LOGGING: bool = (
False # field(default_factory=lambda: os.getenv("ENABLE_MODEL_VERIFICATION_LOGGING", "false").lower() == "true")
)
@@ -397,6 +396,8 @@ class Settings:
@classmethod
def setup_logger(cls):
handler = colorlog.StreamHandler()
# Suppress APScheduler logs below error
colorlog.getLogger("apscheduler").setLevel(colorlog.ERROR)
colorlog.getLogger("apscheduler.scheduler").setLevel(colorlog.ERROR)
@@ -429,8 +430,6 @@ class Settings:
colorlog.getLogger("ddtrace.writer").setLevel(colorlog.INFO)
colorlog.getLogger("ddtrace.internal").setLevel(colorlog.INFO)
colorlog.getLogger("ddtrace.internal.module").setLevel(colorlog.INFO)
colorlog.getLogger("ddtrace.internal.telemetry").setLevel(colorlog.INFO)
colorlog.getLogger("ddtrace.internal.telemetry.writer").setLevel(colorlog.INFO)
colorlog.getLogger("ddtrace.internal.runtime").setLevel(colorlog.INFO)
colorlog.getLogger("ddtrace.internal.runtime.container").setLevel(colorlog.INFO)
colorlog.getLogger("ddtrace._trace.processor").setLevel(colorlog.INFO)
@@ -440,7 +439,10 @@ class Settings:
colorlog.getLogger("datadog").setLevel(colorlog.INFO)
colorlog.getLogger("datadog.dogstatsd").setLevel(colorlog.INFO)
handler = colorlog.StreamHandler()
colorlog.getLogger("ddtrace.internal.telemetry").setLevel(colorlog.INFO)
colorlog.getLogger("ddtrace.internal.telemetry.writer").setLevel(colorlog.INFO)
from pythonjsonlogger.json import JsonFormatter
json_formatter = JsonFormatter(fmt="%(asctime)s %(name)s %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
handler.setFormatter(json_formatter)
apps/pi/pi/services/actions/category_selector.py
-1
View File
@@ -57,5 +57,4 @@ class CategorySelector:
if any(keyword in intent_lower for keyword in ["page", "note", "wiki", "notepad"]):
return "pages"
# No default fallback - let the LLM handle ambiguous cases with NO_ACTIONS_PLANNED
return ""
apps/pi/pi/services/actions/plane_sdk_adapter.py
+1277 -1206
View File
File diff suppressed because it is too large Load Diff
apps/pi/pi/services/actions/prompt_followup.py
+33 -21
View File
@@ -32,6 +32,7 @@ from pi.app.models.enums import MessageMetaStepType
from pi.services.actions.method_executor import MethodExecutor
from pi.services.actions.plane_actions_executor import PlaneActionsExecutor
from pi.services.chat.chat import PlaneChatBot
from pi.services.chat.prompts import plane_context
from pi.services.retrievers.pg_store.action_artifact import add_query_to_artifact
from pi.services.retrievers.pg_store.action_artifact import get_artifact_prompt_history_from_flow_steps
@@ -122,7 +123,7 @@ class EntityResolver:
return tool_results
def _extract_entity_info(self, tool_name: str, tool_args: Dict[str, Any], result: str) -> Dict[str, Any]:
def _extract_entity_info(self, tool_name: str, tool_args: Dict[str, Any], result: Any) -> Dict[str, Any]:
"""Extract entity information using standardized mapping."""
entity_type = TOOL_ENTITY_MAPPING.get(tool_name)
if not entity_type:
@@ -147,8 +148,28 @@ class EntityResolver:
return {result_key: entity_id, name_key: entity_name}
def _extract_entity_id_from_result(self, result: str) -> Optional[str]:
def _extract_entity_id_from_result(self, result: Any) -> Optional[str]:
"""Extract entity ID from tool result using existing patterns."""
# Handle structured dict payload (new format) - highest priority
if isinstance(result, dict):
# Check entity field first
if "entity" in result and isinstance(result["entity"], dict):
entity = result["entity"]
if "entity_id" in entity:
return str(entity["entity_id"])
# Check data field
if "data" in result and isinstance(result["data"], dict):
data = result["data"]
if "id" in data:
return str(data["id"])
# Check if result itself has entity_id
if "entity_id" in result:
return str(result["entity_id"])
# Check if result itself has id
if "id" in result:
return str(result["id"])
# Handle string results (legacy format)
result_str = str(result)
# Look for UUID patterns (standard Plane format)
@@ -314,6 +335,8 @@ class ArtifactFollowupService:
"""Execute LLM with search tools for entity resolution."""
try:
# Bind tools to LLM (same pattern as main flow)
# deduplicate search tools
search_tools = list({tool.name: tool for tool in search_tools}.values())
llm_with_tools = self.chatbot.tool_llm.bind_tools(search_tools)
llm_with_tools.set_tracking_context(message_id, db, MessageMetaStepType.ARTIFACT_MODIFICATION)
@@ -455,22 +478,13 @@ class ArtifactFollowupService:
)
ARTIFACT_MODIFICATION_PROMPT = """# Plane Artifact Modification Assistant
ARTIFACT_MODIFICATION_PROMPT = f"""# Plane Artifact Modification Assistant
You are a specialized AI assistant for modifying Plane project management artifacts.
Your role is to process user requests and generate precise JSON modifications.
## System Context
Plane is an enterprise project management platform with these core entities:
- **Workspaces**: Top-level organizational containers
- **Projects**: Collections of work items, cycles, and modules
- **Work Items**: Core units of work with customizable properties
- **States**: Workflow stages (Backlog, Unstarted, Started, Completed, Cancelled)
- **Cycles**: Time-boxed development sprints
- **Modules**: Logical feature groupings
- **Labels**: Categorization tags
- **Users**: Team members who can be assigned to work
## Context about Plane
{plane_context}
## Your Task
@@ -539,18 +553,17 @@ When user references "me", "myself", "assign to me", or "assign it to me":
After executing any necessary search tools, provide a complete JSON object with all modifications:
```json
{
"field_name": "value",
{{"field_name": "value",
"relation_field_ids": ["uuid1", "uuid2"],
"single_relation_id": "uuid"
}
}}
```
## Examples
**Simple Assignment:**
Request: "assign it to me"
Response: `{"assignee_ids": ["actual-user-uuid-here"]}` (NOT "me")
Response: `{{"assignee_ids": ["actual-user-uuid-here"]}}` (NOT "me")
**Mixed Operations:**
Request: "set state to done, priority to high, assign to me, add dev label"
@@ -560,12 +573,11 @@ Process:
3. Use current user ID for self-assignment
4. Response:
```json
{
"state_id": "state-uuid",
{{"state_id": "state-uuid",
"priority": "high",
"assignee_ids": ["actual-user-uuid-here"],
"label_ids": ["label-uuid"]
}
}}
```
## Critical Requirements
apps/pi/pi/services/actions/registry.py
+15 -11
View File
@@ -44,7 +44,7 @@ API_METHODS: Dict[str, Dict[str, str]] = {
"projects": {
"create": "Create a new project",
"list": "List projects in workspace",
# "retrieve": "Get a single project by ID",
"retrieve": "Get details of a single project",
"update": "Update project details",
"delete": "Delete a project",
"archive": "Archive a project",
@@ -96,14 +96,6 @@ API_METHODS: Dict[str, Dict[str, str]] = {
"create_project_page": "Create a new page in a project",
"create_workspace_page": "Create a new page in the workspace",
},
"assets": {
"create": "Create a new asset",
"create_user_upload": "Upload user-specific assets",
"get_generic": "Retrieve generic assets",
"update_generic": "Update generic assets",
"update_user": "Update user assets",
"delete_user": "Delete user assets",
},
"users": {
"get_current": "Get current user information",
},
@@ -222,7 +214,7 @@ METHOD_NAME_MAP: Dict[str, Dict[str, str]] = {
"add_work_items": "add_cycle_work_items",
"list_work_items": "list_cycle_work_items",
"retrieve_work_item": "retrieve_cycle_work_item",
"remove_work_item": "delete_cycle_work_item",
"remove_work_item": "remove_cycle_work_item",
"transfer_work_items": "transfer_cycle_work_items",
"delete": "delete_cycle",
},
@@ -250,7 +242,7 @@ METHOD_NAME_MAP: Dict[str, Dict[str, str]] = {
"list_archived": "list_archived_modules",
"add_work_items": "add_module_work_items",
"list_work_items": "list_module_work_items",
"remove_work_item": "delete_module_work_item",
"remove_work_item": "remove_module_work_item",
"delete": "delete_module",
},
"pages": {
@@ -335,6 +327,18 @@ METHOD_NAME_MAP: Dict[str, Dict[str, str]] = {
}
# tool_name to category mapping
def build_tool_name_to_category_map() -> Dict[str, str]:
tool_name_to_category_map = {}
for category, methods in METHOD_NAME_MAP.items():
for method_name in methods.values():
tool_name_to_category_map[method_name] = category
return tool_name_to_category_map
TOOL_NAME_TO_CATEGORY_MAP = build_tool_name_to_category_map()
def get_method_name_map(category: str) -> Dict[str, str]:
"""Return the simplified→actual method mapping for a category (empty if none)."""
return METHOD_NAME_MAP.get(category, {}).copy()
apps/pi/pi/services/actions/tools/__init__.py
+18
View File
@@ -35,23 +35,41 @@ from .worklogs import get_worklog_tools
CATEGORY_TO_PROVIDER: Dict[str, Callable] = {
"activity": get_activity_tools,
"activities": get_activity_tools,
"assets": get_asset_tools,
"asset": get_asset_tools,
"attachments": get_attachment_tools,
"attachment": get_attachment_tools,
"comments": get_comment_tools,
"comment": get_comment_tools,
"cycles": get_cycle_tools,
"cycle": get_cycle_tools,
"intake": get_intake_tools,
"intakes": get_intake_tools,
"labels": get_label_tools,
"label": get_label_tools,
"links": get_link_tools,
"link": get_link_tools,
"members": get_member_tools,
"member": get_member_tools,
"modules": get_module_tools,
"module": get_module_tools,
"pages": get_page_tools,
"page": get_page_tools,
"projects": get_project_tools,
"project": get_project_tools,
"properties": get_property_tools,
"property": get_property_tools,
"states": get_state_tools,
"state": get_state_tools,
"types": get_type_tools,
"type": get_type_tools,
"users": get_user_tools,
"user": get_user_tools,
"workitems": get_workitem_tools,
"workitem": get_workitem_tools,
"worklogs": get_worklog_tools,
"worklog": get_worklog_tools,
}
apps/pi/pi/services/actions/tools/activity.py
+8 -6
View File
@@ -2,6 +2,8 @@
Activity API tools for Plane activity tracking operations.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -16,7 +18,7 @@ def get_activity_tools(method_executor, context):
"""Return LangChain tools for the activity category using method_executor and context."""
@tool
async def activity_list(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def activity_list(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List activity for a project.
Args:
@@ -31,16 +33,16 @@ def get_activity_tools(method_executor, context):
result = await method_executor.execute("activity", "list", project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved activity list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved activity list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list activity", result["error"])
return PlaneToolBase.format_error_payload("Failed to list activity", result["error"])
@tool
async def activity_retrieve(
activity_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Get a single activity by ID."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -56,8 +58,8 @@ def get_activity_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved activity", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved activity", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve activity", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve activity", result["error"])
return [activity_list, activity_retrieve]
apps/pi/pi/services/actions/tools/assets.py
+22 -18
View File
@@ -2,6 +2,8 @@
Assets API tools for Plane asset management operations.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -16,7 +18,7 @@ def get_asset_tools(method_executor, context):
"""Return LangChain tools for the assets category using method_executor and context."""
@tool
async def assets_create(project_id: Optional[str] = None, workspace_slug: Optional[str] = None, **kwargs) -> str:
async def assets_create(project_id: Optional[str] = None, workspace_slug: Optional[str] = None, **kwargs) -> Dict[str, Any]:
"""Create a new generic asset.
Args:
@@ -31,12 +33,12 @@ def get_asset_tools(method_executor, context):
result = await method_executor.execute("assets", "create", project_id=project_id, workspace_slug=workspace_slug, **kwargs)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully created asset", result["data"])
return PlaneToolBase.format_success_payload("Successfully created asset", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create asset", result["error"])
return PlaneToolBase.format_error_payload("Failed to create asset", result["error"])
@tool
async def assets_create_user_upload(project_id: Optional[str] = None, workspace_slug: Optional[str] = None, **kwargs) -> str:
async def assets_create_user_upload(project_id: Optional[str] = None, workspace_slug: Optional[str] = None, **kwargs) -> Dict[str, Any]:
"""Upload user-specific assets."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -46,16 +48,16 @@ def get_asset_tools(method_executor, context):
result = await method_executor.execute("assets", "create_user_upload", project_id=project_id, workspace_slug=workspace_slug, **kwargs)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully created user asset upload", result["data"])
return PlaneToolBase.format_success_payload("Successfully created user asset upload", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create user asset upload", result["error"])
return PlaneToolBase.format_error_payload("Failed to create user asset upload", result["error"])
@tool
async def assets_get_generic(
asset_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Retrieve generic assets."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -65,12 +67,14 @@ def get_asset_tools(method_executor, context):
result = await method_executor.execute("assets", "get_generic", asset_id=asset_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved generic asset", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved generic asset", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to get generic asset", result["error"])
return PlaneToolBase.format_error_payload("Failed to get generic asset", result["error"])
@tool
async def assets_update_generic(asset_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None, **kwargs) -> str:
async def assets_update_generic(
asset_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None, **kwargs
) -> Dict[str, Any]:
"""Update generic assets."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -82,12 +86,12 @@ def get_asset_tools(method_executor, context):
"assets", "update_generic", asset_id=asset_id, project_id=project_id, workspace_slug=workspace_slug, **kwargs
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated generic asset", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated generic asset", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update generic asset", result["error"])
return PlaneToolBase.format_error_payload("Failed to update generic asset", result["error"])
@tool
async def assets_update_user(asset_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None, **kwargs) -> str:
async def assets_update_user(asset_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None, **kwargs) -> Dict[str, Any]:
"""Update user assets."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -99,16 +103,16 @@ def get_asset_tools(method_executor, context):
"assets", "update_user", asset_id=asset_id, project_id=project_id, workspace_slug=workspace_slug, **kwargs
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated user asset", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated user asset", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update user asset", result["error"])
return PlaneToolBase.format_error_payload("Failed to update user asset", result["error"])
@tool
async def assets_delete_user(
asset_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Delete user assets."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -118,8 +122,8 @@ def get_asset_tools(method_executor, context):
result = await method_executor.execute("assets", "delete_user", asset_id=asset_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted user asset", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted user asset", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete user asset", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete user asset", result["error"])
return [assets_create, assets_create_user_upload, assets_get_generic, assets_update_generic, assets_update_user, assets_delete_user]
apps/pi/pi/services/actions/tools/attachments.py
+98 -20
View File
@@ -2,6 +2,8 @@
Attachments API tools for Plane file attachment operations.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -19,16 +21,26 @@ def get_attachment_tools(method_executor, context):
async def attachments_create(
issue_id: str,
asset: str,
name: str,
size: int,
file_type: Optional[str] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new attachment on work item.
Args:
issue_id: Parameter description (required)
asset: Parameter description (required)
project_id: Parameter description (optional)
workspace_slug: Parameter description (optional)
issue_id: UUID of the work item to attach file to (required)
asset: Asset identifier or URL (required)
name: Original filename of the asset (required)
size: File size in bytes (required)
file_type: MIME type of the file (optional)
external_id: External identifier for the asset (optional)
external_source: External source system (optional)
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -41,17 +53,28 @@ def get_attachment_tools(method_executor, context):
"create",
issue_id=issue_id,
asset=asset,
name=name,
size=size,
type=file_type,
external_id=external_id,
external_source=external_source,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully created attachment", result["data"])
return PlaneToolBase.format_success_payload("Successfully created attachment", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create attachment", result["error"])
return PlaneToolBase.format_error_payload("Failed to create attachment", result["error"])
@tool
async def attachments_list(issue_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
"""List attachments for an issue."""
async def attachments_list(issue_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List attachments for a work item.
Args:
issue_id: UUID of the work item (required)
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -60,17 +83,25 @@ def get_attachment_tools(method_executor, context):
result = await method_executor.execute("attachments", "list", issue_id=issue_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved attachments list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved attachments list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list attachments", result["error"])
return PlaneToolBase.format_error_payload("Failed to list attachments", result["error"])
@tool
async def attachments_retrieve(
attachment_id: str,
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Get a single attachment by ID."""
) -> Dict[str, Any]:
"""Get a single attachment by ID.
Args:
attachment_id: UUID of the attachment (required)
issue_id: UUID of the work item (required)
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -81,21 +112,67 @@ def get_attachment_tools(method_executor, context):
"attachments",
"retrieve",
attachment_id=attachment_id,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved attachment", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved attachment", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve attachment", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve attachment", result["error"])
@tool
async def attachments_update(
attachment_id: str,
issue_id: str,
is_uploaded: Optional[bool] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> Dict[str, Any]:
"""Update an attachment for a work item.
Args:
attachment_id: UUID of the attachment (required)
issue_id: UUID of the work item (required)
is_uploaded: Mark attachment as uploaded (optional)
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
if project_id is None and "project_id" in context:
project_id = context["project_id"]
result = await method_executor.execute(
"attachments",
"update",
attachment_id=attachment_id,
issue_id=issue_id,
is_uploaded=is_uploaded,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_payload("Successfully updated attachment", result["data"])
else:
return PlaneToolBase.format_error_payload("Failed to update attachment", result["error"])
@tool
async def attachments_delete(
attachment_id: str,
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Delete an attachment."""
) -> Dict[str, Any]:
"""Delete an attachment.
Args:
attachment_id: UUID of the attachment (required)
issue_id: UUID of the work item (required)
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -106,12 +183,13 @@ def get_attachment_tools(method_executor, context):
"attachments",
"delete",
attachment_id=attachment_id,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted attachment", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted attachment", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete attachment", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete attachment", result["error"])
return [attachments_create, attachments_list, attachments_retrieve, attachments_delete]
return [attachments_create, attachments_list, attachments_retrieve, attachments_update, attachments_delete]
apps/pi/pi/services/actions/tools/base.py
+90 -4
View File
@@ -6,6 +6,7 @@ import time
from typing import Any
from typing import Dict
from typing import Optional
from typing import Tuple
from pi import logger
@@ -41,6 +42,25 @@ class PlaneToolBase:
"""Format successful operation response."""
return f"✅ {message}\n\nResult: {data}"
@staticmethod
def format_success_payload(
message: str, data: Any = None, entity: Optional[Dict[str, Any]] = None, meta: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
"""
Structured success payload for API/consumers. Intended long-term replacement for string responses.
"""
payload: Dict[str, Any] = {
"ok": True,
"message": f"✅ {message}",
}
if entity is not None:
payload["entity"] = entity
if data is not None:
payload["data"] = data
if meta is not None:
payload["meta"] = meta
return payload
@staticmethod
async def format_success_response_with_url(message: str, data: Any, entity_type: str, context: Dict[str, Any]) -> str:
"""
@@ -127,11 +147,75 @@ class PlaneToolBase:
# Final fallback to basic response if entity extraction completely fails
return PlaneToolBase.format_success_response(message, data)
@staticmethod
async def format_success_payload_with_url(message: str, data: Any, entity_type: str, context: Dict[str, Any]) -> Dict[str, Any]:
"""
Structured success payload that also includes constructed entity URL information when available.
"""
# Lazy import to avoid circular dependency
from pi.agents.sql_agent.tools import extract_entity_from_api_response
from pi.config import settings as _settings
entity: Dict[str, Any] = {}
try:
entity_data = extract_entity_from_api_response(data, entity_type)
if entity_data:
workspace_slug = context.get("workspace_slug")
frontend_url = _settings.plane_api.FRONTEND_URL
if context.get("workspace_id"):
entity_data["workspace"] = str(context["workspace_id"])
if workspace_slug:
try:
from pi.agents.sql_agent.tools import construct_action_entity_url
url_info = await construct_action_entity_url(entity_data, entity_type, workspace_slug, frontend_url)
if url_info:
# Normalize keys to entity_* for consumers
entity["entity_url"] = url_info.get("entity_url")
entity["entity_name"] = url_info.get("entity_name")
entity["entity_type"] = url_info.get("entity_type") or entity_type
entity["entity_id"] = url_info.get("entity_id")
# Carry identifier fields where applicable
if url_info.get("issue_identifier"):
entity["issue_identifier"] = url_info["issue_identifier"]
if url_info.get("entity_identifier"):
entity["entity_identifier"] = url_info["entity_identifier"]
except Exception as e:
log.error(f"Error constructing entity URL: {e}")
# Ensure id/name/type present even if URL creation failed
if entity_data.get("id") and "entity_id" not in entity:
entity["entity_id"] = str(entity_data["id"])
if entity_data.get("name") and "entity_name" not in entity:
entity["entity_name"] = str(entity_data["name"])
if "entity_type" not in entity:
entity["entity_type"] = entity_type
except Exception as e:
log.warning(f"Failed to build structured entity payload: {e}")
return PlaneToolBase.format_success_payload(message, data=data, entity=entity or None)
@staticmethod
def format_error_response(message: str, error: Any) -> str:
"""Format error response."""
return f"❌ {message}\n\nError: {error}"
@staticmethod
def format_error_payload(message: str, error: Any = None, data: Any = None, meta: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
"""
Structured error payload for API/consumers. Intended long-term replacement for string responses.
"""
payload: Dict[str, Any] = {
"ok": False,
"message": f"❌ {message}",
}
if error is not None:
payload["error"] = error
if data is not None:
payload["data"] = data
if meta is not None:
payload["meta"] = meta
return payload
@staticmethod
def generate_project_identifier(name: str) -> str:
"""Generate a project identifier from name."""
@@ -143,10 +227,12 @@ class PlaneToolBase:
return base_identifier
@staticmethod
def generate_fallback_identifier(base_identifier: str) -> str:
"""Generate fallback identifier with timestamp."""
timestamp = str(int(time.time()))[-3:] # Last 3 digits of timestamp
return f"{base_identifier}{timestamp}"
def generate_fallback_name_identifier(base_name: str, base_identifier: str) -> Tuple[str, str]:
"""Generate fallback name and identifier with timestamp."""
timestamp = str(int(time.time()))[-5:] # Last 5 digits of timestamp
name = f"{base_name}{timestamp}"
identifier = f"{base_identifier}{timestamp}"
return name, identifier
def get_workspace_slug_from_context(context: Dict[str, Any]) -> Optional[str]:
apps/pi/pi/services/actions/tools/comments.py
+66 -23
View File
@@ -2,6 +2,8 @@
Comments API tools for Plane issue comment management.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -21,15 +23,20 @@ def get_comment_tools(method_executor, context):
comment_html: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
external_source: Optional[str] = None,
external_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Create a comment on an issue.
Args:
issue_id: Parameter description (required)
comment_html: Parameter description (required)
project_id: Parameter description (optional)
workspace_slug: Parameter description (optional)
issue_id: Issue ID (required)
comment_html: Comment content in HTML format (required)
project_id: Project ID (optional, auto-filled from context)
workspace_slug: Workspace slug (optional, auto-filled from context)
external_source: External source identifier (e.g., "jira")
external_id: External system ID
"""
print(f"Context in comments_create: {context}")
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -41,16 +48,18 @@ def get_comment_tools(method_executor, context):
"create",
issue_id=issue_id,
comment_html=comment_html,
external_source=external_source,
external_id=external_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully created comment", result["data"])
return PlaneToolBase.format_success_payload("Successfully created comment", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create comment", result["error"])
return PlaneToolBase.format_error_payload("Failed to create comment", result["error"])
@tool
async def comments_list(issue_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def comments_list(issue_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List comments for an issue."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -60,17 +69,25 @@ def get_comment_tools(method_executor, context):
result = await method_executor.execute("comments", "list", issue_id=issue_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved comments list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved comments list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list comments", result["error"])
return PlaneToolBase.format_error_payload("Failed to list comments", result["error"])
@tool
async def comments_retrieve(
comment_id: str,
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Get a single comment by ID."""
) -> Dict[str, Any]:
"""Get a single comment by ID.
Args:
comment_id: Comment ID (required)
issue_id: Issue ID (required)
project_id: Project ID (optional, auto-filled from context)
workspace_slug: Workspace slug (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -81,22 +98,36 @@ def get_comment_tools(method_executor, context):
"comments",
"retrieve",
comment_id=comment_id,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved comment", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved comment", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve comment", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve comment", result["error"])
@tool
async def comments_update(
comment_id: str,
comment_html: str,
issue_id: str,
comment_html: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Update comment details."""
external_source: Optional[str] = None,
external_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Update comment details.
Args:
comment_id: Comment ID (required)
issue_id: Issue ID (required)
comment_html: Comment content in HTML format
project_id: Project ID (optional, auto-filled from context)
workspace_slug: Workspace slug (optional, auto-filled from context)
external_source: External source identifier (e.g., "jira")
external_id: External system ID
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -107,22 +138,33 @@ def get_comment_tools(method_executor, context):
"comments",
"update",
comment_id=comment_id,
issue_id=issue_id,
comment_html=comment_html,
external_source=external_source,
external_id=external_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated comment", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated comment", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update comment", result["error"])
return PlaneToolBase.format_error_payload("Failed to update comment", result["error"])
@tool
async def comments_delete(
comment_id: str,
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Delete a comment."""
) -> Dict[str, Any]:
"""Delete a comment.
Args:
comment_id: Comment ID (required)
issue_id: Issue ID (required)
project_id: Project ID (optional, auto-filled from context)
workspace_slug: Workspace slug (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -133,12 +175,13 @@ def get_comment_tools(method_executor, context):
"comments",
"delete",
comment_id=comment_id,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted comment", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted comment", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete comment", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete comment", result["error"])
return [comments_create, comments_list, comments_retrieve, comments_update, comments_delete]
apps/pi/pi/services/actions/tools/cycles.py
+55 -59
View File
@@ -2,6 +2,8 @@
Cycles API tools for Plane project cycle management operations.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -26,19 +28,22 @@ def get_cycle_tools(method_executor, context):
owned_by: Optional[str] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new cycle.
Args:
name: Cycle name (required)
project_id: Project ID (required - provide from conversation context or previous actions)
workspace_slug: Workspace slug (required - provide from conversation context)
start_date: Start date in YYYY-MM-DD format (optional)
end_date: End date in YYYY-MM-DD format (optional)
start_date: Start date in YYYY-MM-DD format (IMPORTANT: if provided, end_date must also be provided)
end_date: End date in YYYY-MM-DD format (IMPORTANT: if provided, start_date must also be provided)
description: Cycle description (optional)
owned_by: User ID who owns the cycle (provide when user specifies ownership or assignment)
external_id: External system identifier (optional)
external_source: External system source name (optional)
Note: Both start_date and end_date must be provided together, or neither should be provided.
You cannot provide only one date - the API will reject the request.
"""
# Auto-fill from context if not provided
if workspace_slug is None:
@@ -59,15 +64,14 @@ def get_cycle_tools(method_executor, context):
end_date=end_date,
description=description,
owned_by=owned_by,
user_id=context.get("user_id"),
external_id=external_id,
external_source=external_source,
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url(f"Successfully created cycle '{name}'", result["data"], "cycle", context)
return await PlaneToolBase.format_success_payload_with_url(f"Successfully created cycle '{name}'", result["data"], "cycle", context)
else:
return PlaneToolBase.format_error_response("Failed to create cycle", result["error"])
return PlaneToolBase.format_error_payload("Failed to create cycle", result["error"])
@tool
async def cycles_list(
@@ -76,7 +80,7 @@ def get_cycle_tools(method_executor, context):
per_page: Optional[int] = 20,
cursor: Optional[str] = None,
cycle_view: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""List cycles in a project.
Args:
@@ -100,7 +104,7 @@ def get_cycle_tools(method_executor, context):
# Validate required parameters after auto-fill
if not project_id or not workspace_slug:
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"Failed to list cycles",
"Missing required context: project_id/workspace_slug",
)
@@ -116,31 +120,31 @@ def get_cycle_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved cycles list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved cycles list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list cycles", result["error"])
return PlaneToolBase.format_error_payload("Failed to list cycles", result["error"])
@tool
async def cycles_retrieve(pk: str, project_id: str, workspace_slug: str) -> str:
async def cycles_retrieve(cycle_id: str, project_id: str, workspace_slug: str) -> Dict[str, Any]:
"""Retrieve a single cycle by ID.
Args:
pk: Cycle ID (required)
cycle_id: Cycle ID (required)
project_id: Project ID (required - provide from conversation context or previous actions)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
"""
# project_id and workspace_slug are now required parameters
result = await method_executor.execute("cycles", "retrieve", pk=pk, project_id=project_id, workspace_slug=workspace_slug)
result = await method_executor.execute("cycles", "retrieve", cycle_id=cycle_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved cycle", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved cycle", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve cycle", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve cycle", result["error"])
@tool
async def cycles_update(
pk: str,
cycle_id: str,
project_id: str,
workspace_slug: Optional[str] = None,
name: Optional[str] = None,
@@ -150,11 +154,11 @@ def get_cycle_tools(method_executor, context):
owned_by: Optional[str] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Update cycle details.
Args:
pk: Cycle ID (required)
cycle_id: Cycle ID (required)
project_id: Project ID (required - provide from conversation context or previous actions)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
name: New cycle name
@@ -186,15 +190,17 @@ def get_cycle_tools(method_executor, context):
if external_source is not None:
update_data["external_source"] = external_source
result = await method_executor.execute("cycles", "update", pk=pk, project_id=project_id, workspace_slug=workspace_slug, **update_data)
result = await method_executor.execute(
"cycles", "update", cycle_id=cycle_id, project_id=project_id, workspace_slug=workspace_slug, **update_data
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url("Successfully updated cycle", result["data"], "cycle", context)
return await PlaneToolBase.format_success_payload_with_url("Successfully updated cycle", result["data"], "cycle", context)
else:
return PlaneToolBase.format_error_response("Failed to update cycle", result["error"])
return PlaneToolBase.format_error_payload("Failed to update cycle", result["error"])
@tool
async def cycles_archive(cycle_id: str, project_id: str, workspace_slug: str) -> str:
async def cycles_archive(cycle_id: str, project_id: str, workspace_slug: str) -> Dict[str, Any]:
"""Archive a cycle.
Args:
@@ -204,15 +210,15 @@ def get_cycle_tools(method_executor, context):
"""
# project_id and workspace_slug are now required parameters
result = await method_executor.execute("cycles", "archive", pk=cycle_id, project_id=project_id, workspace_slug=workspace_slug)
result = await method_executor.execute("cycles", "archive", cycle_id=cycle_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully archived cycle", result["data"])
return PlaneToolBase.format_success_payload("Successfully archived cycle", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to archive cycle", result["error"])
return PlaneToolBase.format_error_payload("Failed to archive cycle", result["error"])
@tool
async def cycles_add_work_items(cycle_id: str, issues: list, project_id: str, workspace_slug: Optional[str] = None) -> str:
async def cycles_add_work_items(cycle_id: str, issues: list, project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Add work items to a cycle.
Args:
@@ -227,7 +233,7 @@ def get_cycle_tools(method_executor, context):
workspace_slug = context["workspace_slug"]
# Check if issues list is empty
if not issues or len(issues) == 0:
return PlaneToolBase.format_success_response(
return PlaneToolBase.format_success_payload(
"No work items to add to cycle", {"message": "There are no work items satisfying the criteria to add to the cycle", "issues_count": 0}
)
@@ -237,12 +243,12 @@ def get_cycle_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully added work items to cycle", result["data"])
return PlaneToolBase.format_success_payload("Successfully added work items to cycle", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to add work items to cycle", result["error"])
return PlaneToolBase.format_error_payload("Failed to add work items to cycle", result["error"])
@tool
async def cycles_remove_work_item(cycle_id: str, issue_id: str, project_id: str, workspace_slug: Optional[str] = None) -> str:
async def cycles_remove_work_item(cycle_id: str, issue_id: str, project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Remove a work item from a cycle.
Args:
@@ -262,12 +268,12 @@ def get_cycle_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully removed work item from cycle", result["data"])
return PlaneToolBase.format_success_payload("Successfully removed work item from cycle", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to remove work item from cycle", result["error"])
return PlaneToolBase.format_error_payload("Failed to remove work item from cycle", result["error"])
@tool
async def cycles_unarchive(cycle_id: str, project_id: str, workspace_slug: Optional[str] = None) -> str:
async def cycles_unarchive(cycle_id: str, project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Unarchive a cycle."""
# Auto-fill from context if not provided
if workspace_slug is None:
@@ -275,14 +281,14 @@ def get_cycle_tools(method_executor, context):
workspace_slug = context["workspace_slug"]
# project_id and workspace_slug are now required parameters
result = await method_executor.execute("cycles", "unarchive", pk=cycle_id, project_id=project_id, workspace_slug=workspace_slug)
result = await method_executor.execute("cycles", "unarchive", cycle_id=cycle_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response(f"Successfully unarchived cycle '{cycle_id}'", result["data"])
return PlaneToolBase.format_success_payload(f"Successfully unarchived cycle '{cycle_id}'", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to unarchive cycle", result["error"])
return PlaneToolBase.format_error_payload("Failed to unarchive cycle", result["error"])
@tool
async def cycles_list_archived(project_id: str, workspace_slug: Optional[str] = None) -> str:
async def cycles_list_archived(project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List archived cycles."""
# Auto-fill from context if not provided
if workspace_slug is None:
@@ -292,12 +298,12 @@ def get_cycle_tools(method_executor, context):
# project_id and workspace_slug are now required parameters
result = await method_executor.execute("cycles", "list_archived", project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved archived cycles list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved archived cycles list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list archived cycles", result["error"])
return PlaneToolBase.format_error_payload("Failed to list archived cycles", result["error"])
@tool
async def cycles_list_work_items(cycle_id: str, project_id: str, workspace_slug: Optional[str] = None) -> str:
async def cycles_list_work_items(cycle_id: str, project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List work items in a cycle."""
# Auto-fill from context if not provided
if workspace_slug is None:
@@ -307,12 +313,12 @@ def get_cycle_tools(method_executor, context):
# project_id and workspace_slug are now required parameters
result = await method_executor.execute("cycles", "list_work_items", cycle_id=cycle_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved cycle work items list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved cycle work items list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list cycle work items", result["error"])
return PlaneToolBase.format_error_payload("Failed to list cycle work items", result["error"])
@tool
async def cycles_retrieve_work_item(cycle_id: str, issue_id: str, project_id: str, workspace_slug: Optional[str] = None) -> str:
async def cycles_retrieve_work_item(cycle_id: str, issue_id: str, project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Get a specific work item in a cycle."""
# Auto-fill from context if not provided
if workspace_slug is None:
@@ -324,12 +330,12 @@ def get_cycle_tools(method_executor, context):
"cycles", "retrieve_work_item", cycle_id=cycle_id, issue_id=issue_id, project_id=project_id, workspace_slug=workspace_slug
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved cycle work item", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved cycle work item", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve cycle work item", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve cycle work item", result["error"])
@tool
async def cycles_transfer_work_items(cycle_id: str, new_cycle_id: str, project_id: str, workspace_slug: Optional[str] = None) -> str:
async def cycles_transfer_work_items(cycle_id: str, new_cycle_id: str, project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Transfer work items between cycles."""
# Auto-fill from context if not provided
if workspace_slug is None:
@@ -341,23 +347,13 @@ def get_cycle_tools(method_executor, context):
"cycles", "transfer_work_items", cycle_id=cycle_id, new_cycle_id=new_cycle_id, project_id=project_id, workspace_slug=workspace_slug
)
if result["success"]:
return PlaneToolBase.format_success_response(
return PlaneToolBase.format_success_payload(
f"Successfully transferred work items from cycle '{cycle_id}' to '{new_cycle_id}'", result["data"]
)
else:
return PlaneToolBase.format_error_response("Failed to transfer cycle work items", result["error"])
return PlaneToolBase.format_error_payload("Failed to transfer cycle work items", result["error"])
# Get entity search tools relevant only to cycles
from .entity_search import get_entity_search_tools
entity_search_tools = get_entity_search_tools(method_executor, context)
cycle_entity_search_tools = [
t
for t in entity_search_tools
if getattr(t, "name", "").find("cycle") != -1 or getattr(t, "name", "").find("user") != -1 or getattr(t, "name", "") == "search_current_cycle"
]
return cycle_entity_search_tools + [
return [
cycles_create,
cycles_update,
cycles_archive,
apps/pi/pi/services/actions/tools/entity_search.py
+92 -77
View File
@@ -4,10 +4,13 @@ Provides search functionality for all major entities using database queries.
"""
import uuid
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
from pi import settings
from pi.core.db.plane import PlaneDBPool
from .base import PlaneToolBase
@@ -42,7 +45,7 @@ def get_entity_search_tools(method_executor, context):
return str(row["id"]) if row else None
@tool
async def search_module_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def search_module_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Search for a module by name and return its ID.
Args:
@@ -64,17 +67,17 @@ def get_entity_search_tools(method_executor, context):
result = await search_module_by_name(name, project_id, workspace_slug)
if result:
return PlaneToolBase.format_success_response(
return PlaneToolBase.format_success_payload(
f"Found module '{name}'", {"id": result["id"], "name": result["name"], "project_id": result["project_id"]}
)
else:
return PlaneToolBase.format_error_response(f"No module found with name '{name}'", "Not found")
return PlaneToolBase.format_success_payload(f"No module found with name '{name}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for module '{name}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for module '{name}': {str(e)}", str(e))
@tool
async def list_member_projects(workspace_id: Optional[str] = None, user_id: Optional[str] = None, limit: Optional[int] = 50) -> str:
async def list_member_projects(workspace_id: Optional[str] = None, user_id: Optional[str] = None, limit: Optional[int] = 50) -> Dict[str, Any]:
"""List active projects the current user is a member of (archived/deleted excluded).
Args:
@@ -90,7 +93,7 @@ def get_entity_search_tools(method_executor, context):
try:
if not workspace_id or not user_id:
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"Failed to list projects",
"Missing workspace_id/user_id in context",
)
@@ -109,7 +112,7 @@ def get_entity_search_tools(method_executor, context):
ORDER BY p.name
LIMIT $3
"""
rows = await PlaneDBPool.fetch(query, (workspace_id, user_id, int(limit or 50)))
rows = await PlaneDBPool.fetch(query, (uuid.UUID(workspace_id), uuid.UUID(user_id), int(limit or 50)))
projects = [
{
"id": str(r["id"]),
@@ -119,12 +122,12 @@ def get_entity_search_tools(method_executor, context):
}
for r in rows
]
return PlaneToolBase.format_success_response("Successfully retrieved member projects", {"projects": projects, "count": len(projects)})
return PlaneToolBase.format_success_payload("Successfully retrieved member projects", {"projects": projects, "count": len(projects)})
except Exception as e:
return PlaneToolBase.format_error_response("Failed to list member projects", str(e))
return PlaneToolBase.format_error_payload("Failed to list member projects", str(e))
@tool
async def search_project_by_name(name: str, workspace_slug: Optional[str] = None, user_id: Optional[str] = None) -> str:
async def search_project_by_name(name: str, workspace_slug: Optional[str] = None, user_id: Optional[str] = None) -> Dict[str, Any]:
"""Search for a project by name.
Args:
@@ -167,15 +170,15 @@ def get_entity_search_tools(method_executor, context):
else:
message = f"Found 1 project matching '{name}'"
return PlaneToolBase.format_success_response(message, data)
return PlaneToolBase.format_success_payload(message, data)
else:
return PlaneToolBase.format_error_response(f"No project found with name '{name}'", "Not found")
return PlaneToolBase.format_success_payload(f"No project found with name '{name}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for project '{name}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for project '{name}': {str(e)}", str(e))
@tool
async def search_project_by_identifier(identifier: str, workspace_slug: Optional[str] = None) -> str:
async def search_project_by_identifier(identifier: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Search for a project by its identifier (e.g., 'HYDR', 'PARM') and return its UUID.
CRITICAL: Project identifiers are short uppercase codes (like 'HYDR', 'PARM'), NOT UUIDs.
@@ -212,7 +215,7 @@ def get_entity_search_tools(method_executor, context):
result = await PlaneDBPool.fetchrow(query, tuple(params))
if result:
return PlaneToolBase.format_success_response(
return PlaneToolBase.format_success_payload(
f"Found project with identifier '{identifier}'",
{
"id": str(result["id"]),
@@ -222,13 +225,13 @@ def get_entity_search_tools(method_executor, context):
},
)
else:
return PlaneToolBase.format_error_response(f"No project found with identifier '{identifier}'", "Not found")
return PlaneToolBase.format_success_payload(f"No project found with identifier '{identifier}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for project '{identifier}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for project '{identifier}': {str(e)}", str(e))
@tool
async def search_cycle_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def search_cycle_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Search for a cycle by name and return its ID.
Args:
@@ -255,15 +258,15 @@ def get_entity_search_tools(method_executor, context):
response_data["start_date"] = result["start_date"]
if result.get("end_date"):
response_data["end_date"] = result["end_date"]
return PlaneToolBase.format_success_response(f"Found cycle '{name}'", response_data)
return PlaneToolBase.format_success_payload(f"Found cycle '{name}'", response_data)
else:
return PlaneToolBase.format_error_response(f"No cycle found with name '{name}'", "Not found")
return PlaneToolBase.format_success_payload(f"No cycle found with name '{name}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for cycle '{name}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for cycle '{name}': {str(e)}", str(e))
@tool
async def search_current_cycle(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def search_current_cycle(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Search for the current active cycle in a project (where today's date falls within the cycle's start and end dates).
IMPORTANT: Use this tool when the user asks about the "current cycle", "active cycle", "ongoing cycle", or "this cycle".
@@ -281,7 +284,7 @@ def get_entity_search_tools(method_executor, context):
try:
if not project_id:
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"Failed to search for current cycle",
"project_id is required to search for current cycle",
)
@@ -289,52 +292,64 @@ def get_entity_search_tools(method_executor, context):
# Normalize project_id if an identifier like 'OGX' was passed
normalized_project_id = await _normalize_project_id(project_id, workspace_slug)
if not normalized_project_id:
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"Failed to search for current cycle",
"project_id is required to search for current cycle",
)
from pi.app.api.v1.helpers.plane_sql_queries import search_current_cycle
result = await search_current_cycle(normalized_project_id, workspace_slug)
results = await search_current_cycle(normalized_project_id, workspace_slug)
if result:
response_data = {
"id": result["id"],
"name": result["name"],
"project_id": result["project_id"],
"start_date": result.get("start_date"),
"end_date": result.get("end_date"),
"is_current": True,
}
# Build cycle URL and include in response message
cycle_url = None
if results:
cycles_data = []
cycle_names = []
try:
from pi import settings
ws_slug = result.get("workspace_slug") or workspace_slug
if ws_slug and result.get("project_id") and result.get("id"):
api_base_url = settings.plane_api.FRONTEND_URL
cycle_url = f"{api_base_url}/{ws_slug}/projects/{result["project_id"]}/cycles/{result["id"]}/"
response_data["url"] = cycle_url
api_base_url = settings.plane_api.FRONTEND_URL
except Exception:
pass
api_base_url = None
# Include URL in the message text so LLM sees it even without entity_urls
message = f"Found current active cycle: '{result["name"]}'"
if cycle_url:
message += f"\nCycle URL: {cycle_url}"
for result in results:
cycle_data = {
"id": result["id"],
"name": result["name"],
"project_id": result["project_id"],
"start_date": result.get("start_date"),
"end_date": result.get("end_date"),
"is_current": True,
}
return PlaneToolBase.format_success_response(message, response_data)
# Build cycle URL and include in response message
cycle_url = None
if api_base_url:
ws_slug = result.get("workspace_slug") or workspace_slug
if ws_slug and result.get("project_id") and result.get("id"):
cycle_url = f"{api_base_url}/{ws_slug}/projects/{result["project_id"]}/cycles/{result["id"]}/"
cycle_data["url"] = cycle_url
cycles_data.append(cycle_data)
cycle_names.append(result["name"])
# Build response message
if len(results) == 1:
message = f"Found current active cycle: '{cycle_names[0]}'"
if cycles_data[0].get("url"):
message += f"\nCycle URL: {cycles_data[0]["url"]}"
response_data = cycles_data[0]
else:
cycle_names_str = ", ".join([f"'{name}'" for name in cycle_names])
message = f"Found {len(results)} current active cycles: {cycle_names_str}"
response_data = {"cycles": cycles_data, "count": len(results)}
return PlaneToolBase.format_success_payload(message, response_data)
else:
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_success_payload(
"No current active cycle found",
"No cycle is active today (no cycle where today falls between start_date and end_date)",
)
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for current cycle: {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for current cycle: {str(e)}", str(e))
@tool
async def list_recent_cycles(
@@ -342,7 +357,7 @@ def get_entity_search_tools(method_executor, context):
status: Optional[str] = "completed",
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""List recent cycles in the current project/workspace context.
Use this to resolve "last cycle" or "previous cycle" in project chats.
@@ -370,7 +385,7 @@ def get_entity_search_tools(method_executor, context):
workspace_id = str(row["id"]) if row else None
if not project_id and not workspace_id:
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"Failed to list recent cycles",
"Missing project context; provide project_id or ensure project chat context",
)
@@ -388,12 +403,12 @@ def get_entity_search_tools(method_executor, context):
if rows:
data["last_cycle_id"] = rows[0]["id"]
return PlaneToolBase.format_success_response("Successfully retrieved recent cycles", data)
return PlaneToolBase.format_success_payload("Successfully retrieved recent cycles", data)
except Exception as e:
return PlaneToolBase.format_error_response("Failed to list recent cycles", str(e))
return PlaneToolBase.format_error_payload("Failed to list recent cycles", str(e))
@tool
async def search_label_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def search_label_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Search for a label by name and return its ID.
Args:
@@ -415,17 +430,17 @@ def get_entity_search_tools(method_executor, context):
result = await search_label_by_name(name, project_id, workspace_slug)
if result:
return PlaneToolBase.format_success_response(
return PlaneToolBase.format_success_payload(
f"Found label '{name}'", {"id": result["id"], "name": result["name"], "project_id": result["project_id"]}
)
else:
return PlaneToolBase.format_error_response(f"No label found with name '{name}'", "Not found")
return PlaneToolBase.format_success_payload(f"No label found with name '{name}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for label '{name}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for label '{name}': {str(e)}", str(e))
@tool
async def search_state_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def search_state_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Search for a state by name and return its ID.
Args:
@@ -447,14 +462,14 @@ def get_entity_search_tools(method_executor, context):
result = await search_state_by_name(name, project_id, workspace_slug)
if result:
return PlaneToolBase.format_success_response(
return PlaneToolBase.format_success_payload(
f"Found state '{name}'", {"id": result["id"], "name": result["name"], "project_id": result["project_id"]}
)
else:
return PlaneToolBase.format_error_response(f"No state found with name '{name}'", "Not found")
return PlaneToolBase.format_success_payload(f"No state found with name '{name}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for state '{name}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for state '{name}': {str(e)}", str(e))
@tool
async def search_user_by_name(
@@ -462,7 +477,7 @@ def get_entity_search_tools(method_executor, context):
workspace_slug: Optional[str] = None,
first_name: Optional[str] = None,
last_name: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Search for a user by display, first, or last name and return matches.
Args:
@@ -478,7 +493,7 @@ def get_entity_search_tools(method_executor, context):
# Ensure at least one name parameter is present
if not any([display_name, first_name, last_name]):
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"At least one of display_name, first_name, or last_name must be provided",
"missing_parameters",
)
@@ -535,15 +550,15 @@ def get_entity_search_tools(method_executor, context):
else:
message = f"Found 1 user matching '{display_name}'"
return PlaneToolBase.format_success_response(message, users_data)
return PlaneToolBase.format_success_payload(message, users_data)
else:
return PlaneToolBase.format_error_response(f"No user found with display name '{display_name}'", "Not found")
return PlaneToolBase.format_success_payload(f"No user found with display name '{display_name}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for user '{display_name}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for user '{display_name}': {str(e)}", str(e))
@tool
async def search_workitem_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def search_workitem_by_name(name: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Search for a work item by name and return its ID.
Args:
@@ -565,7 +580,7 @@ def get_entity_search_tools(method_executor, context):
result = await search_workitem_by_name(name, project_id, workspace_slug)
if result:
return PlaneToolBase.format_success_response(
return PlaneToolBase.format_success_payload(
f"Found work item '{name}'",
{
"id": result["id"],
@@ -577,13 +592,13 @@ def get_entity_search_tools(method_executor, context):
},
)
else:
return PlaneToolBase.format_error_response(f"No work item found with name '{name}'", "Not found")
return PlaneToolBase.format_success_payload(f"No work item found with name '{name}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for work item '{name}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for work item '{name}': {str(e)}", str(e))
@tool
async def search_workitem_by_identifier(identifier: str, workspace_slug: Optional[str] = None) -> str:
async def search_workitem_by_identifier(identifier: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Search for a work item by its unique identifier (e.g., 'WEB-821') and return its details.
Don't call this tool if the 'identifier' string is not in the format 'PROJECT-SEQUENCE'.
@@ -601,7 +616,7 @@ def get_entity_search_tools(method_executor, context):
result = await search_workitem_by_identifier(identifier, workspace_slug)
if result:
return PlaneToolBase.format_success_response(
return PlaneToolBase.format_success_payload(
f"Found work item '{identifier}'",
{
"id": result["id"],
@@ -619,10 +634,10 @@ def get_entity_search_tools(method_executor, context):
},
)
else:
return PlaneToolBase.format_error_response(f"No work item found with identifier '{identifier}'", "Not found")
return PlaneToolBase.format_success_payload(f"No work item found with identifier '{identifier}'", "Not found")
except Exception as e:
return PlaneToolBase.format_error_response(f"Error searching for work item '{identifier}': {str(e)}", str(e))
return PlaneToolBase.format_error_payload(f"Error searching for work item '{identifier}': {str(e)}", str(e))
return [
list_member_projects,
apps/pi/pi/services/actions/tools/intake.py
+15 -15
View File
@@ -28,7 +28,7 @@ def get_intake_tools(method_executor, context):
assignee: Optional[str] = None,
reporter: Optional[str] = None,
labels: Optional[List[str]] = None,
) -> str:
) -> Dict[str, Any]:
"""Submit work item to intake queue for triage.
Args:
@@ -61,9 +61,9 @@ def get_intake_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response(f"Successfully submitted intake item '{name}'", result["data"])
return PlaneToolBase.format_success_payload(f"Successfully submitted intake item '{name}'", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create intake item", result["error"])
return PlaneToolBase.format_error_payload("Failed to create intake item", result["error"])
@tool
async def intake_list(
@@ -71,7 +71,7 @@ def get_intake_tools(method_executor, context):
workspace_slug: Optional[str] = None,
per_page: Optional[int] = 20,
cursor: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""List intake items awaiting triage."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -89,16 +89,16 @@ def get_intake_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved intake list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved intake list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list intake items", result["error"])
return PlaneToolBase.format_error_payload("Failed to list intake items", result["error"])
@tool
async def intake_retrieve(
intake_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Get a single intake work item by ID."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -115,9 +115,9 @@ def get_intake_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved intake item", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved intake item", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve intake item", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve intake item", result["error"])
@tool
async def intake_update(
@@ -130,7 +130,7 @@ def get_intake_tools(method_executor, context):
assignee: Optional[str] = None,
reporter: Optional[str] = None,
labels: Optional[List[str]] = None,
) -> str:
) -> Dict[str, Any]:
"""Update intake work item details."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -163,16 +163,16 @@ def get_intake_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated intake item", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated intake item", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update intake item", result["error"])
return PlaneToolBase.format_error_payload("Failed to update intake item", result["error"])
@tool
async def intake_delete(
intake_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Remove intake work item."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -189,8 +189,8 @@ def get_intake_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted intake item", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted intake item", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete intake item", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete intake item", result["error"])
return [intake_create, intake_list, intake_retrieve, intake_update, intake_delete]
apps/pi/pi/services/actions/tools/labels.py
+57 -23
View File
@@ -2,6 +2,8 @@
Labels API tools for Plane labeling operations.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -22,8 +24,24 @@ def get_label_tools(method_executor, context):
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
description: Optional[str] = None,
) -> str:
"""Create a new label."""
parent: Optional[str] = None,
sort_order: Optional[float] = None,
external_source: Optional[str] = None,
external_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Create a new label.
Args:
name: Label name (required)
color: Label color in hex format (required)
project_id: Project ID (required - provide from conversation context or previous actions)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
description: Label description
parent: Parent label ID for nested labels
sort_order: Sort order for display (float)
external_source: External source identifier (e.g., "jira")
external_id: External system ID
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -38,18 +56,22 @@ def get_label_tools(method_executor, context):
project_id=project_id,
workspace_slug=workspace_slug,
description=description,
parent=parent,
sort_order=sort_order,
external_source=external_source,
external_id=external_id,
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url(f"Successfully created label '{name}'", result["data"], "label", context)
return await PlaneToolBase.format_success_payload_with_url(f"Successfully created label '{name}'", result["data"], "label", context)
else:
return PlaneToolBase.format_error_response("Failed to create label", result["error"])
return PlaneToolBase.format_error_payload("Failed to create label", result["error"])
@tool
async def labels_list(
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""List all labels in a project."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -65,16 +87,16 @@ def get_label_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved labels list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved labels list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list labels", result["error"])
return PlaneToolBase.format_error_payload("Failed to list labels", result["error"])
@tool
async def labels_retrieve(
label_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Retrieve details of a specific label.
Args:
@@ -97,9 +119,9 @@ def get_label_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved label details", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved label details", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve label", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve label", result["error"])
@tool
async def labels_update(
@@ -109,7 +131,11 @@ def get_label_tools(method_executor, context):
name: Optional[str] = None,
color: Optional[str] = None,
description: Optional[str] = None,
) -> str:
parent: Optional[str] = None,
sort_order: Optional[float] = None,
external_source: Optional[str] = None,
external_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Update an existing label.
Args:
@@ -117,6 +143,10 @@ def get_label_tools(method_executor, context):
name: New label name
color: New label color
description: New label description
parent: Parent label ID for nested labels
sort_order: Sort order for display (float)
external_source: External source identifier (e.g., "jira")
external_id: External system ID
project_id: Project ID (required - provide from conversation context or previous actions)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
"""
@@ -127,7 +157,19 @@ def get_label_tools(method_executor, context):
project_id = context["project_id"]
# Build update data with only non-None values
update_data = {k: v for k, v in {"name": name, "color": color, "description": description}.items() if v is not None}
update_data = {
k: v
for k, v in {
"name": name,
"color": color,
"description": description,
"parent": parent,
"sort_order": sort_order,
"external_source": external_source,
"external_id": external_id,
}.items()
if v is not None
}
result = await method_executor.execute(
"labels",
@@ -139,16 +181,8 @@ def get_label_tools(method_executor, context):
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url(f"Successfully updated label '{label_id}'", result["data"], "label", context)
return await PlaneToolBase.format_success_payload_with_url(f"Successfully updated label '{label_id}'", result["data"], "label", context)
else:
return PlaneToolBase.format_error_response("Failed to update label", result["error"])
return PlaneToolBase.format_error_payload("Failed to update label", result["error"])
# Get entity search tools relevant only to labels
from .entity_search import get_entity_search_tools
entity_search_tools = get_entity_search_tools(method_executor, context)
label_entity_search_tools = [
t for t in entity_search_tools if getattr(t, "name", "").find("label") != -1 or getattr(t, "name", "").find("user") != -1
]
return [labels_create, labels_list, labels_retrieve, labels_update] + label_entity_search_tools
return [labels_create, labels_list, labels_retrieve, labels_update]
apps/pi/pi/services/actions/tools/links.py
+54 -23
View File
@@ -2,6 +2,8 @@
Links API tools for Plane issue link management.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -22,15 +24,15 @@ def get_link_tools(method_executor, context):
title: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a link for an issue.
Args:
issue_id: Parameter description (required)
url: Parameter description (required)
title: Parameter description (optional)
project_id: Parameter description (optional)
workspace_slug: Parameter description (optional)
issue_id: UUID of the work item to create a link for (required)
url: The URL of the link (required)
title: Optional title/label for the link
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -48,12 +50,12 @@ def get_link_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully created link", result["data"])
return PlaneToolBase.format_success_payload("Successfully created link", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create link", result["error"])
return PlaneToolBase.format_error_payload("Failed to create link", result["error"])
@tool
async def links_list(issue_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def links_list(issue_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List links for an issue."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -63,17 +65,25 @@ def get_link_tools(method_executor, context):
result = await method_executor.execute("links", "list", issue_id=issue_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved links list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved links list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list links", result["error"])
return PlaneToolBase.format_error_payload("Failed to list links", result["error"])
@tool
async def links_retrieve(
link_id: str,
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Get a single link by ID."""
) -> Dict[str, Any]:
"""Get a single link by ID.
Args:
link_id: UUID of the link to retrieve (required)
issue_id: UUID of the work item the link belongs to (required)
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -84,23 +94,34 @@ def get_link_tools(method_executor, context):
"links",
"retrieve",
link_id=link_id,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved link", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved link", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve link", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve link", result["error"])
@tool
async def links_update(
link_id: str,
issue_id: str,
url: Optional[str] = None,
title: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Update link details."""
) -> Dict[str, Any]:
"""Update link details.
Args:
link_id: UUID of the link to update (required)
issue_id: UUID of the work item the link belongs to (required)
url: New URL for the link (optional)
title: New title/label for the link (optional)
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -118,22 +139,31 @@ def get_link_tools(method_executor, context):
"links",
"update",
link_id=link_id,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
**update_data,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated link", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated link", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update link", result["error"])
return PlaneToolBase.format_error_payload("Failed to update link", result["error"])
@tool
async def links_delete(
link_id: str,
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Delete a link."""
) -> Dict[str, Any]:
"""Delete a link.
Args:
link_id: UUID of the link to delete (required)
issue_id: UUID of the work item the link belongs to (required)
project_id: UUID of the project (optional, auto-filled from context)
workspace_slug: Workspace slug identifier (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -144,12 +174,13 @@ def get_link_tools(method_executor, context):
"links",
"delete",
link_id=link_id,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted link", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted link", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete link", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete link", result["error"])
return [links_create, links_list, links_retrieve, links_update, links_delete]
apps/pi/pi/services/actions/tools/members.py
+11 -8
View File
@@ -2,6 +2,8 @@
Members API tools for Plane workspace and project member management.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -17,7 +19,7 @@ def get_member_tools(method_executor, context):
"""Get all Members API tools."""
@tool
async def members_get_workspace_members(workspace_slug: Optional[str] = None) -> str:
async def members_get_workspace_members(workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Get all workspace members (excludes bot users).
Args:
@@ -28,6 +30,7 @@ def get_member_tools(method_executor, context):
workspace_slug = context["workspace_slug"]
result = await method_executor.execute("members", "get_workspace_members", workspace_slug=workspace_slug)
if result["success"]:
# Filter out bot users from the results
members_data = result["data"]
@@ -43,18 +46,18 @@ def get_member_tools(method_executor, context):
or (isinstance(member.get("email"), str) and "[email protected]" in member.get("email", "").lower())
)
]
return PlaneToolBase.format_success_response("Successfully retrieved workspace members", filtered_members)
return PlaneToolBase.format_success_payload("Successfully retrieved workspace members", filtered_members)
else:
# If data is not a list, return as-is (edge case)
return PlaneToolBase.format_success_response("Successfully retrieved workspace members", members_data)
return PlaneToolBase.format_success_payload("Successfully retrieved workspace members", members_data)
else:
return PlaneToolBase.format_error_response("Failed to get workspace members", result["error"])
return PlaneToolBase.format_error_payload("Failed to get workspace members", result["error"])
@tool
async def members_get_project_members(
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Get all project members (excludes bot users).
Args:
@@ -88,11 +91,11 @@ def get_member_tools(method_executor, context):
or (isinstance(member.get("email"), str) and "[email protected]" in member.get("email", "").lower())
)
]
return PlaneToolBase.format_success_response("Successfully retrieved project members", filtered_members)
return PlaneToolBase.format_success_payload("Successfully retrieved project members", filtered_members)
else:
# If data is not a list, return as-is (edge case)
return PlaneToolBase.format_success_response("Successfully retrieved project members", members_data)
return PlaneToolBase.format_success_payload("Successfully retrieved project members", members_data)
else:
return PlaneToolBase.format_error_response("Failed to get project members", result["error"])
return PlaneToolBase.format_error_payload("Failed to get project members", result["error"])
return [members_get_workspace_members, members_get_project_members]
apps/pi/pi/services/actions/tools/modules.py
+40 -44
View File
@@ -35,7 +35,7 @@ def get_module_tools(method_executor, context):
members: Optional[list] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new module.
Args:
@@ -73,12 +73,12 @@ def get_module_tools(method_executor, context):
external_source=external_source,
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url(f"Successfully created module '{name}'", result["data"], "module", context)
return await PlaneToolBase.format_success_payload_with_url(f"Successfully created module '{name}'", result["data"], "module", context)
else:
return PlaneToolBase.format_error_response("Failed to create module", result["error"])
return PlaneToolBase.format_error_payload("Failed to create module", result["error"])
@tool
async def modules_list(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def modules_list(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List all modules in a project.
Args:
@@ -98,19 +98,19 @@ def get_module_tools(method_executor, context):
# Validate project_id is present
if not project_id:
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"Missing required parameter",
"project_id is required to list modules. You MUST first call projects_list to see available projects, OR call ask_for_clarification with disambiguation_options populated by calling projects_list to ask which project the user wants.", # noqa: E501
)
result = await method_executor.execute("modules", "list", project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved modules list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved modules list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list modules", result["error"])
return PlaneToolBase.format_error_payload("Failed to list modules", result["error"])
@tool
async def modules_add_work_items(module_id: str, issues: list, project_id: Optional[str] = None) -> str:
async def modules_add_work_items(module_id: str, issues: list, project_id: Optional[str] = None) -> Dict[str, Any]:
"""Add work items to a module.
Args:
@@ -129,12 +129,12 @@ def get_module_tools(method_executor, context):
"modules", "add_work_items", module_id=module_id, issues=issues, project_id=project_id, workspace_slug=workspace_slug
)
if result["success"]:
return PlaneToolBase.format_success_response(f"Successfully added {len(issues)} work items to module", result["data"])
return PlaneToolBase.format_success_payload(f"Successfully added {len(issues)} work items to module", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to add work items to module", result["error"])
return PlaneToolBase.format_error_payload("Failed to add work items to module", result["error"])
@tool
async def modules_retrieve(module_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def modules_retrieve(module_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Retrieve details of a specific module.
Args:
@@ -148,11 +148,11 @@ def get_module_tools(method_executor, context):
if project_id is None and "project_id" in context:
project_id = context["project_id"]
result = await method_executor.execute("modules", "retrieve", pk=module_id, project_id=project_id, workspace_slug=workspace_slug)
result = await method_executor.execute("modules", "retrieve", module_id=module_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url("Successfully retrieved module details", result["data"], "module", context)
return await PlaneToolBase.format_success_payload_with_url("Successfully retrieved module details", result["data"], "module", context)
else:
return PlaneToolBase.format_error_response("Failed to retrieve module", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve module", result["error"])
@tool
async def modules_update(
@@ -168,7 +168,7 @@ def get_module_tools(method_executor, context):
members: Optional[list] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Update module details.
Args:
@@ -212,15 +212,17 @@ def get_module_tools(method_executor, context):
if external_source is not None:
update_data["external_source"] = external_source
result = await method_executor.execute("modules", "update", pk=module_id, project_id=project_id, workspace_slug=workspace_slug, **update_data)
result = await method_executor.execute(
"modules", "update", module_id=module_id, project_id=project_id, workspace_slug=workspace_slug, **update_data
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url("Successfully updated module", result["data"], "module", context)
return await PlaneToolBase.format_success_payload_with_url("Successfully updated module", result["data"], "module", context)
else:
return PlaneToolBase.format_error_response("Failed to update module", result["error"])
return PlaneToolBase.format_error_payload("Failed to update module", result["error"])
@tool
async def modules_archive(module_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def modules_archive(module_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Archive a module.
Args:
@@ -234,15 +236,15 @@ def get_module_tools(method_executor, context):
if project_id is None and "project_id" in context:
project_id = context["project_id"]
result = await method_executor.execute("modules", "archive", pk=module_id, project_id=project_id, workspace_slug=workspace_slug)
result = await method_executor.execute("modules", "archive", module_id=module_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully archived module", result["data"])
return PlaneToolBase.format_success_payload("Successfully archived module", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to archive module", result["error"])
return PlaneToolBase.format_error_payload("Failed to archive module", result["error"])
@tool
async def modules_unarchive(module_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def modules_unarchive(module_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Unarchive a module.
Args:
@@ -256,15 +258,15 @@ def get_module_tools(method_executor, context):
if project_id is None and "project_id" in context:
project_id = context["project_id"]
result = await method_executor.execute("modules", "unarchive", pk=module_id, project_id=project_id, workspace_slug=workspace_slug)
result = await method_executor.execute("modules", "unarchive", module_id=module_id, project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully unarchived module", result["data"])
return PlaneToolBase.format_success_payload("Successfully unarchived module", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to unarchive module", result["error"])
return PlaneToolBase.format_error_payload("Failed to unarchive module", result["error"])
@tool
async def modules_list_archived(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def modules_list_archived(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List archived modules in a project.
Args:
@@ -279,12 +281,12 @@ def get_module_tools(method_executor, context):
result = await method_executor.execute("modules", "list_archived", project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved archived modules list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved archived modules list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list archived modules", result["error"])
return PlaneToolBase.format_error_payload("Failed to list archived modules", result["error"])
@tool
async def modules_list_work_items(module_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def modules_list_work_items(module_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List work items in a module.
Args:
@@ -302,12 +304,14 @@ def get_module_tools(method_executor, context):
"modules", "list_work_items", module_id=module_id, project_id=project_id, workspace_slug=workspace_slug
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved module work items", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved module work items", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list module work items", result["error"])
return PlaneToolBase.format_error_payload("Failed to list module work items", result["error"])
@tool
async def modules_remove_work_item(module_id: str, issue_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def modules_remove_work_item(
module_id: str, issue_id: str, project_id: Optional[str] = None, workspace_slug: Optional[str] = None
) -> Dict[str, Any]:
"""Remove a work item from a module.
Args:
@@ -327,17 +331,9 @@ def get_module_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully removed work item from module", result["data"])
return PlaneToolBase.format_success_payload("Successfully removed work item from module", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to remove work item from module", result["error"])
# Get entity search tools relevant only to modules
from .entity_search import get_entity_search_tools
entity_search_tools = get_entity_search_tools(method_executor, context)
module_entity_search_tools = [
t for t in entity_search_tools if getattr(t, "name", "").find("module") != -1 or getattr(t, "name", "").find("user") != -1
]
return PlaneToolBase.format_error_payload("Failed to remove work item from module", result["error"])
return [
modules_create,
@@ -350,4 +346,4 @@ def get_module_tools(method_executor, context):
modules_list_archived,
modules_list_work_items,
modules_remove_work_item,
] + module_entity_search_tools
]
apps/pi/pi/services/actions/tools/pages.py
+13 -29
View File
@@ -2,6 +2,8 @@
Pages API tools for Plane documentation and page management operations.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -26,7 +28,7 @@ def get_page_tools(method_executor, context):
access: Optional[int] = None,
color: Optional[str] = None,
logo_props: Optional[dict] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new page in a project.
Args:
@@ -71,10 +73,10 @@ def get_page_tools(method_executor, context):
# Using "project" not "project_id" because extract_entity_from_api_response looks for "project"
if project_id:
data["project"] = project_id
return await PlaneToolBase.format_success_response_with_url(f"Project page '{name}' created successfully", data, "page", context)
return await PlaneToolBase.format_success_payload_with_url(f"Project page '{name}' created successfully", data, "page", context)
else:
error_msg = result.get("error", "Unknown error occurred")
return PlaneToolBase.format_error_response("Failed to create project page", error_msg)
return PlaneToolBase.format_error_payload("Failed to create project page", error_msg)
@tool
async def pages_create_workspace_page(
@@ -83,7 +85,7 @@ def get_page_tools(method_executor, context):
access: Optional[int] = None,
color: Optional[str] = None,
logo_props: Optional[dict] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new page in the workspace.
Args:
@@ -117,10 +119,10 @@ def get_page_tools(method_executor, context):
if result.get("success"):
data = result.get("data", {})
return await PlaneToolBase.format_success_response_with_url(f"Workspace page '{name}' created successfully", data, "page", context)
return await PlaneToolBase.format_success_payload_with_url(f"Workspace page '{name}' created successfully", data, "page", context)
else:
error_msg = result.get("error", "Unknown error occurred")
return PlaneToolBase.format_error_response("Failed to create workspace page", error_msg)
return PlaneToolBase.format_error_payload("Failed to create workspace page", error_msg)
# In workspace/global context, expose a single tool that forces clarification for scope
@tool
@@ -131,7 +133,7 @@ def get_page_tools(method_executor, context):
access: Optional[int] = None,
color: Optional[str] = None,
logo_props: Optional[dict] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new page either in a project or at the workspace level.
Args:
@@ -185,34 +187,16 @@ def get_page_tools(method_executor, context):
# for the URL builder to construct the correct project page URL
if scope_label == "Project" and project_id != "__workspace_scope__":
data["project"] = project_id
return await PlaneToolBase.format_success_response_with_url(f"{scope_label} page '{name}' created successfully", data, "page", context)
return await PlaneToolBase.format_success_payload_with_url(f"{scope_label} page '{name}' created successfully", data, "page", context)
else:
error_msg = result.get("error", "Unknown error occurred")
return PlaneToolBase.format_error_response(f"Failed to create {scope_label.lower()} page", error_msg)
return PlaneToolBase.format_error_payload(f"Failed to create {scope_label.lower()} page", error_msg)
# Return tools based on context:
# - In project chat: only project pages (workspace pages are not accessible in project UI)
# - In workspace/global chat: consolidated tool that requires scope selection via clarification
# CRITICAL: Include ONLY the entity search tools relevant to Pages
# Keep the tool surface minimal to avoid unintended tool selection
entity_tools = []
try:
from .entity_search import get_entity_search_tools
all_entity_tools = get_entity_search_tools(method_executor, context) or []
allowed_entity_tool_names = {
# Projects-only for page scope selection
"list_member_projects",
"search_project_by_name",
"search_project_by_identifier",
}
entity_tools = [t for t in all_entity_tools if getattr(t, "name", "") in allowed_entity_tool_names]
except Exception:
# Best-effort; if entity search tools are unavailable, proceed with page tools only
pass
if is_project_chat:
return [pages_create_project_page] + entity_tools
return [pages_create_project_page]
else:
return [pages_create_page] + entity_tools
return [pages_create_page]
apps/pi/pi/services/actions/tools/projects.py
+101 -35
View File
@@ -2,6 +2,7 @@
Projects API tools for Plane workspace operations.
"""
import re
from typing import Any
from typing import Dict
from typing import Optional
@@ -9,6 +10,7 @@ from typing import Optional
from langchain_core.tools import tool
from pi import logger
from pi.core.db import PlaneDBPool
from .base import PlaneToolBase
@@ -42,7 +44,9 @@ def get_project_tools(method_executor, context):
archive_in: Optional[int] = None,
close_in: Optional[int] = None,
timezone: Optional[str] = None,
) -> str:
time_tracking_enabled: Optional[bool] = None,
is_issue_type_enabled: Optional[bool] = None,
) -> Dict[str, Any]:
"""Create a new project in the workspace.
Args:
@@ -64,6 +68,8 @@ def get_project_tools(method_executor, context):
close_in: Auto-close workitems after N months. Should be less than or equal to 12 (optional)
timezone: Timezone for the project (optional)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
time_tracking_enabled: Enable/disable time tracking, also called worklogs (optional)
is_issue_type_enabled: Enable/disable issue type, also called workitem types (optional)
"""
# Determine identifier
base_identifier = identifier or PlaneToolBase.generate_project_identifier(name)
@@ -110,40 +116,73 @@ def get_project_tools(method_executor, context):
payload["close_in"] = close_in
if timezone is not None:
payload["timezone"] = timezone
if time_tracking_enabled is not None:
payload["time_tracking_enabled"] = time_tracking_enabled
if is_issue_type_enabled is not None:
payload["is_issue_type_enabled"] = is_issue_type_enabled
# Try to create project
result = await method_executor.execute("projects", "create", **payload)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url(
return await PlaneToolBase.format_success_payload_with_url(
f"Successfully created project '{name}' with identifier '{base_identifier}'", result["data"], "project", context
)
else:
# If failed due to identifier conflict, try with a different identifier
if "already taken" in result.get("error", "") or "409" in result.get("error", ""):
# Generate a new identifier with timestamp
new_identifier = PlaneToolBase.generate_fallback_identifier(base_identifier)
new_name, new_identifier = PlaneToolBase.generate_fallback_name_identifier(name, base_identifier)
# Retry with new identifier (preserve other fields)
payload["name"] = new_name
payload["identifier"] = new_identifier
retry_result = await method_executor.execute("projects", "create", **payload)
if retry_result["success"]:
return await PlaneToolBase.format_success_response_with_url(
f"Successfully created project '{name}' with identifier '{new_identifier}' (original '{base_identifier}' was taken)",
return await PlaneToolBase.format_success_payload_with_url(
f"Successfully created project '{new_name}' with identifier '{new_identifier}' (original '{name}' was taken)",
retry_result["data"],
"project",
context,
)
else:
log.info(f"Failed to create project. Error: {retry_result["error"]}\nPayload: {payload}")
return PlaneToolBase.format_error_response("Failed to create project even with alternative identifier", retry_result["error"])
return PlaneToolBase.format_error_payload("Failed to create project even with alternative identifier", retry_result["error"])
else:
# Check if project was actually created despite the error (e.g. timeout)
try:
# We need to check if the project exists in the DB with the same identifier and workspace_slug
query = """
SELECT p.id, p.name, p.identifier, p.workspace_id
FROM projects p
JOIN workspaces w ON p.workspace_id = w.id
WHERE p.identifier = $1 AND w.slug = $2 AND p.deleted_at IS NULL
"""
row = await PlaneDBPool.fetchrow(query, (base_identifier, workspace_slug))
if row:
project_data = {
"id": str(row["id"]),
"name": row["name"],
"identifier": row["identifier"],
"workspace_id": str(row["workspace_id"]),
# Add workspace_slug for context
"workspace_slug": workspace_slug,
}
return await PlaneToolBase.format_success_payload_with_url(
f"Successfully created project '{name}' with identifier '{base_identifier}'",
project_data,
"project",
context,
)
except Exception as e:
log.error(f"Failed to recover project creation from DB: {e}")
log.info(f"Failed to create project. Error: {result["error"]}\nPayload: {payload}")
return PlaneToolBase.format_error_response("Failed to create project", result["error"])
return PlaneToolBase.format_error_payload("Failed to create project", result["error"])
@tool
async def projects_list(workspace_slug: Optional[str] = None, per_page: Optional[int] = 20, cursor: Optional[str] = None) -> str:
async def projects_list(workspace_slug: Optional[str] = None, per_page: Optional[int] = 20, cursor: Optional[str] = None) -> Dict[str, Any]:
"""List projects in the workspace.
Args:
@@ -158,32 +197,52 @@ def get_project_tools(method_executor, context):
result = await method_executor.execute("projects", "list", workspace_slug=workspace_slug, per_page=per_page, cursor=cursor)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved projects list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved projects list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list projects", result["error"])
return PlaneToolBase.format_error_payload("Failed to list projects", result["error"])
@tool
async def projects_retrieve(pk: str, workspace_slug: Optional[str] = None) -> str:
async def projects_retrieve(project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Retrieve a single project by ID.
Args:
pk: Project ID (required)
project_id: Project ID (required)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
"""
# Lightweight validation to prevent bad calls and noisy 404s
# - Reject placeholders like "<id of project: X>"
# - Require a UUID-shaped project_id
try:
if "<id of" in project_id:
return PlaneToolBase.format_error_payload(
"Invalid project_id: received a placeholder. Resolve a real UUID using search_project_by_name or search_project_by_identifier before calling projects_retrieve.", # noqa E501
f"project_id={project_id}",
)
uuid_regex = re.compile(r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")
if not uuid_regex.match(project_id):
return PlaneToolBase.format_error_payload(
"Invalid project_id format: expected UUID. Use search_project_by_name or search_project_by_identifier to resolve the UUID, then retry.", # noqa E501
f"project_id={project_id}",
)
except Exception:
# Best-effort validation; proceed if validation fails unexpectedly
pass
# Auto-fill workspace_slug from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
result = await method_executor.execute("projects", "retrieve", pk=pk, workspace_slug=workspace_slug)
result = await method_executor.execute("projects", "retrieve", project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved project", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved project", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve project", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve project", result["error"])
@tool
async def projects_update(
pk: str,
project_id: str,
workspace_slug: Optional[str] = None,
name: Optional[str] = None,
description: Optional[str] = None,
@@ -198,15 +257,17 @@ def get_project_tools(method_executor, context):
issue_views_view: Optional[bool] = None,
page_view: Optional[bool] = None,
intake_view: Optional[bool] = None,
is_time_tracking_enabled: Optional[bool] = None,
is_issue_type_enabled: Optional[bool] = None,
guest_view_all_features: Optional[bool] = None,
archive_in: Optional[int] = None,
close_in: Optional[int] = None,
timezone: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Update project details.
Args:
pk: Project ID (required)
project_id: Project ID (required)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
name: New project name
description: New project description
@@ -225,6 +286,8 @@ def get_project_tools(method_executor, context):
archive_in: Auto-archive issues after N days (optional)
close_in: Auto-close issues after N days (optional)
timezone: Timezone for the project (optional)
is_time_tracking_enabled: Enable/disable time tracking, also called worklogs (optional)
is_issue_type_enabled: Enable/disable issue type, also called workitem types (optional)
"""
# Auto-fill workspace_slug from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -258,6 +321,10 @@ def get_project_tools(method_executor, context):
update_data["page_view"] = page_view
if intake_view is not None:
update_data["intake_view"] = intake_view
if is_time_tracking_enabled is not None:
update_data["is_time_tracking_enabled"] = is_time_tracking_enabled
if is_issue_type_enabled is not None:
update_data["is_issue_type_enabled"] = is_issue_type_enabled
if guest_view_all_features is not None:
update_data["guest_view_all_features"] = guest_view_all_features
if archive_in is not None:
@@ -270,18 +337,18 @@ def get_project_tools(method_executor, context):
result = await method_executor.execute(
"projects",
"update",
pk=pk,
project_id=project_id,
workspace_slug=workspace_slug,
**update_data,
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url("Successfully updated project", result["data"], "project", context)
return await PlaneToolBase.format_success_payload_with_url("Successfully updated project", result["data"], "project", context)
else:
return PlaneToolBase.format_error_response("Failed to update project", result["error"])
return PlaneToolBase.format_error_payload("Failed to update project", result["error"])
@tool
async def projects_archive(project_id: str, workspace_slug: Optional[str] = None) -> str:
async def projects_archive(project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Archive a project.
Args:
@@ -295,12 +362,12 @@ def get_project_tools(method_executor, context):
result = await method_executor.execute("projects", "archive", project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully archived project", result["data"])
return PlaneToolBase.format_success_payload("Successfully archived project", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to archive project", result["error"])
return PlaneToolBase.format_error_payload("Failed to archive project", result["error"])
@tool
async def projects_unarchive(project_id: str, workspace_slug: Optional[str] = None) -> str:
async def projects_unarchive(project_id: str, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""Restore an archived project.
Args:
@@ -314,16 +381,15 @@ def get_project_tools(method_executor, context):
result = await method_executor.execute("projects", "unarchive", project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully unarchived project", result["data"])
return PlaneToolBase.format_success_payload("Successfully unarchived project", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to unarchive project", result["error"])
return PlaneToolBase.format_error_payload("Failed to unarchive project", result["error"])
# Get entity search tools relevant only to projects
from .entity_search import get_entity_search_tools
entity_search_tools = get_entity_search_tools(method_executor, context)
project_entity_search_tools = [
t for t in entity_search_tools if getattr(t, "name", "").find("project") != -1 or getattr(t, "name", "").find("user") != -1
return [
projects_create,
projects_list,
projects_retrieve,
projects_update,
projects_archive,
projects_unarchive,
]
return [projects_create, projects_list, projects_retrieve, projects_update, projects_archive, projects_unarchive] + project_entity_search_tools
apps/pi/pi/services/actions/tools/properties.py
+170 -53
View File
@@ -2,6 +2,8 @@
Properties API tools for Plane custom property management.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -18,20 +20,40 @@ def get_property_tools(method_executor, context):
@tool
async def properties_create(
name: str,
display_name: str,
property_type: str,
type_id: str,
relation_type: Optional[str] = None,
description: Optional[str] = None,
is_required: Optional[bool] = None,
default_value: Optional[list] = None,
settings: Optional[dict] = None,
is_active: Optional[bool] = None,
is_multi: Optional[bool] = None,
validation_rules: Optional[dict] = None,
external_source: Optional[str] = None,
external_id: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Create a new property.
) -> Dict[str, Any]:
"""Create a new work item property.
Args:
name: Parameter description (required)
type_id: Parameter description (required)
description: Parameter description (optional)
project_id: Parameter description (optional)
workspace_slug: Parameter description (optional)
display_name: Display name for the property (required)
property_type: Type of property - TEXT, DATETIME, DECIMAL, BOOLEAN, OPTION, RELATION, URL, EMAIL, FILE (required)
type_id: Work item type ID this property belongs to (required)
relation_type: Relation type for RELATION properties - ONE_TO_ONE, ONE_TO_MANY, MANY_TO_ONE, MANY_TO_MANY
description: Description of the property
is_required: Whether this property is required for work items
default_value: List of default values for the property
settings: Property settings (TextAttributeSettings for TEXT, DateAttributeSettings for DATETIME)
is_active: Whether the property is active
is_multi: Whether this property supports multiple values
validation_rules: Validation rules for the property
external_source: External system source identifier (e.g., 'github', 'jira')
external_id: External system's identifier for this property
project_id: Project ID (auto-filled from context if not provided)
workspace_slug: Workspace slug (auto-filled from context if not provided)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -42,39 +64,73 @@ def get_property_tools(method_executor, context):
result = await method_executor.execute(
"properties",
"create",
name=name,
display_name=display_name,
property_type=property_type,
type_id=type_id,
relation_type=relation_type,
description=description,
is_required=is_required,
default_value=default_value,
settings=settings,
is_active=is_active,
is_multi=is_multi,
validation_rules=validation_rules,
external_source=external_source,
external_id=external_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response(f"Successfully created property '{name}'", result["data"])
return PlaneToolBase.format_success_payload(f"Successfully created property '{display_name}'", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create property", result["error"])
return PlaneToolBase.format_error_payload("Failed to create property", result["error"])
@tool
async def properties_list(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
"""List all custom properties."""
async def properties_list(
type_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> Dict[str, Any]:
"""List all work item properties for a specific work item type.
Args:
type_id: Work item type ID to list properties for (required)
project_id: Project ID (auto-filled from context if not provided)
workspace_slug: Workspace slug (auto-filled from context if not provided)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
if project_id is None and "project_id" in context:
project_id = context["project_id"]
result = await method_executor.execute("properties", "list", project_id=project_id, workspace_slug=workspace_slug)
result = await method_executor.execute(
"properties",
"list",
type_id=type_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved properties list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved properties list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list properties", result["error"])
return PlaneToolBase.format_error_payload("Failed to list properties", result["error"])
@tool
async def properties_retrieve(
property_id: str,
type_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Get a single property by ID."""
) -> Dict[str, Any]:
"""Get a single work item property by ID.
Args:
property_id: Property ID to retrieve (required)
type_id: Work item type ID this property belongs to (required)
project_id: Project ID (auto-filled from context if not provided)
workspace_slug: Workspace slug (auto-filled from context if not provided)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -85,23 +141,54 @@ def get_property_tools(method_executor, context):
"properties",
"retrieve",
property_id=property_id,
type_id=type_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved property", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved property", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve property", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve property", result["error"])
@tool
async def properties_update(
property_id: str,
name: Optional[str] = None,
type_id: str,
display_name: Optional[str] = None,
property_type: Optional[str] = None,
relation_type: Optional[str] = None,
description: Optional[str] = None,
is_required: Optional[bool] = None,
default_value: Optional[list] = None,
settings: Optional[dict] = None,
is_active: Optional[bool] = None,
is_multi: Optional[bool] = None,
validation_rules: Optional[dict] = None,
external_source: Optional[str] = None,
external_id: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Update property details."""
) -> Dict[str, Any]:
"""Update work item property details.
Args:
property_id: Property ID to update (required)
type_id: Work item type ID this property belongs to (required)
display_name: Display name for the property
property_type: Type of property - TEXT, DATETIME, DECIMAL, BOOLEAN, OPTION, RELATION, URL, EMAIL, FILE
relation_type: Relation type for RELATION properties - ONE_TO_ONE, ONE_TO_MANY, MANY_TO_ONE, MANY_TO_MANY
description: Description of the property
is_required: Whether this property is required for work items
default_value: List of default values for the property
settings: Property settings (TextAttributeSettings for TEXT, DateAttributeSettings for DATETIME)
is_active: Whether the property is active
is_multi: Whether this property supports multiple values
validation_rules: Validation rules for the property
external_source: External system source identifier (e.g., 'github', 'jira')
external_id: External system's identifier for this property
project_id: Project ID (auto-filled from context if not provided)
workspace_slug: Workspace slug (auto-filled from context if not provided)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -109,32 +196,61 @@ def get_property_tools(method_executor, context):
project_id = context["project_id"]
# Build update data
update_data = {}
if name is not None:
update_data["name"] = name
update_data: dict[str, Any] = {}
if display_name is not None:
update_data["display_name"] = display_name
if property_type is not None:
update_data["property_type"] = property_type
if relation_type is not None:
update_data["relation_type"] = relation_type
if description is not None:
update_data["description"] = description
if is_required is not None:
update_data["is_required"] = is_required
if default_value is not None:
update_data["default_value"] = default_value
if settings is not None:
update_data["settings"] = settings
if is_active is not None:
update_data["is_active"] = is_active
if is_multi is not None:
update_data["is_multi"] = is_multi
if validation_rules is not None:
update_data["validation_rules"] = validation_rules
if external_source is not None:
update_data["external_source"] = external_source
if external_id is not None:
update_data["external_id"] = external_id
result = await method_executor.execute(
"properties",
"update",
property_id=property_id,
type_id=type_id,
project_id=project_id,
workspace_slug=workspace_slug,
**update_data,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated property", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated property", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update property", result["error"])
return PlaneToolBase.format_error_payload("Failed to update property", result["error"])
@tool
async def properties_delete(
property_id: str,
type_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Delete a property."""
) -> Dict[str, Any]:
"""Delete a work item property.
Args:
property_id: Property ID to delete (required)
type_id: Work item type ID this property belongs to (required)
project_id: Project ID (auto-filled from context if not provided)
workspace_slug: Workspace slug (auto-filled from context if not provided)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -145,13 +261,14 @@ def get_property_tools(method_executor, context):
"properties",
"delete",
property_id=property_id,
type_id=type_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted property", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted property", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete property", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete property", result["error"])
@tool
async def properties_create_option(
@@ -161,7 +278,7 @@ def get_property_tools(method_executor, context):
value: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a property option."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -180,9 +297,9 @@ def get_property_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully created property option", result["data"])
return PlaneToolBase.format_success_payload("Successfully created property option", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create property option", result["error"])
return PlaneToolBase.format_error_payload("Failed to create property option", result["error"])
@tool
async def properties_create_value(
@@ -192,7 +309,7 @@ def get_property_tools(method_executor, context):
value: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a property value."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -211,9 +328,9 @@ def get_property_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully created property value", result["data"])
return PlaneToolBase.format_success_payload("Successfully created property value", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create property value", result["error"])
return PlaneToolBase.format_error_payload("Failed to create property value", result["error"])
@tool
async def properties_list_options(
@@ -221,7 +338,7 @@ def get_property_tools(method_executor, context):
type_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""List property options."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -238,9 +355,9 @@ def get_property_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved property options", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved property options", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list property options", result["error"])
return PlaneToolBase.format_error_payload("Failed to list property options", result["error"])
@tool
async def properties_list_values(
@@ -248,7 +365,7 @@ def get_property_tools(method_executor, context):
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""List property values."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -265,9 +382,9 @@ def get_property_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved property values", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved property values", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list property values", result["error"])
return PlaneToolBase.format_error_payload("Failed to list property values", result["error"])
@tool
async def properties_retrieve_option(
@@ -276,7 +393,7 @@ def get_property_tools(method_executor, context):
option_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Get property option by ID."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -294,9 +411,9 @@ def get_property_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved property option", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved property option", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve property option", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve property option", result["error"])
@tool
async def properties_update_option(
@@ -307,7 +424,7 @@ def get_property_tools(method_executor, context):
value: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Update property option."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -333,9 +450,9 @@ def get_property_tools(method_executor, context):
**update_data,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated property option", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated property option", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update property option", result["error"])
return PlaneToolBase.format_error_payload("Failed to update property option", result["error"])
@tool
async def properties_delete_option(
@@ -344,7 +461,7 @@ def get_property_tools(method_executor, context):
option_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Delete property option."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -362,9 +479,9 @@ def get_property_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted property option", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted property option", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete property option", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete property option", result["error"])
return [
properties_create,
apps/pi/pi/services/actions/tools/states.py
+57 -22
View File
@@ -2,6 +2,8 @@
States API tools for Plane workflow state management.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -23,7 +25,12 @@ def get_state_tools(method_executor, context):
workspace_slug: Optional[str] = None,
description: Optional[str] = None,
group: Optional[str] = None,
) -> str:
sequence: Optional[int] = None,
is_triage: Optional[bool] = None,
default: Optional[bool] = None,
external_source: Optional[str] = None,
external_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Create a new workflow state.
Args:
@@ -33,6 +40,11 @@ def get_state_tools(method_executor, context):
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
description: State description
group: State group (backlog, unstarted, started, completed, cancelled)
sequence: Display sequence order
is_triage: Whether this is a triage state
default: Whether this is the default state
external_source: External source identifier (e.g., "jira")
external_id: External system ID
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -49,18 +61,23 @@ def get_state_tools(method_executor, context):
workspace_slug=workspace_slug,
description=description,
group=group,
sequence=sequence,
is_triage=is_triage,
default=default,
external_source=external_source,
external_id=external_id,
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url(f"Successfully created state '{name}'", result["data"], "state", context)
return await PlaneToolBase.format_success_payload_with_url(f"Successfully created state '{name}'", result["data"], "state", context)
else:
return PlaneToolBase.format_error_response("Failed to create state", result["error"])
return PlaneToolBase.format_error_payload("Failed to create state", result["error"])
@tool
async def states_list(
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""List all workflow states in a project."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -76,16 +93,16 @@ def get_state_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved states list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved states list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list states", result["error"])
return PlaneToolBase.format_error_payload("Failed to list states", result["error"])
@tool
async def states_retrieve(
state_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Retrieve details of a specific workflow state.
Args:
@@ -108,9 +125,9 @@ def get_state_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved state details", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved state details", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve state", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve state", result["error"])
@tool
async def states_update(
@@ -120,7 +137,13 @@ def get_state_tools(method_executor, context):
name: Optional[str] = None,
color: Optional[str] = None,
description: Optional[str] = None,
) -> str:
sequence: Optional[int] = None,
group: Optional[str] = None,
is_triage: Optional[bool] = None,
default: Optional[bool] = None,
external_source: Optional[str] = None,
external_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Update an existing workflow state.
Args:
@@ -128,6 +151,12 @@ def get_state_tools(method_executor, context):
name: New state name
color: New state color
description: New state description
sequence: Display sequence order
group: State group (backlog, unstarted, started, completed, cancelled)
is_triage: Whether this is a triage state
default: Whether this is the default state
external_source: External source identifier (e.g., "jira")
external_id: External system ID
project_id: Project ID (required - provide from conversation context or previous actions)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
"""
@@ -138,7 +167,21 @@ def get_state_tools(method_executor, context):
project_id = context["project_id"]
# Build update data with only non-None values
update_data = {k: v for k, v in {"name": name, "color": color, "description": description}.items() if v is not None}
update_data = {
k: v
for k, v in {
"name": name,
"color": color,
"description": description,
"sequence": sequence,
"group": group,
"is_triage": is_triage,
"default": default,
"external_source": external_source,
"external_id": external_id,
}.items()
if v is not None
}
result = await method_executor.execute(
"states",
@@ -150,16 +193,8 @@ def get_state_tools(method_executor, context):
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url(f"Successfully updated state '{state_id}'", result["data"], "state", context)
return await PlaneToolBase.format_success_payload_with_url(f"Successfully updated state '{state_id}'", result["data"], "state", context)
else:
return PlaneToolBase.format_error_response("Failed to update state", result["error"])
return PlaneToolBase.format_error_payload("Failed to update state", result["error"])
# Get entity search tools relevant only to states
from .entity_search import get_entity_search_tools
entity_search_tools = get_entity_search_tools(method_executor, context)
state_entity_search_tools = [
t for t in entity_search_tools if getattr(t, "name", "").find("state") != -1 or getattr(t, "name", "").find("user") != -1
]
return [states_create, states_list, states_retrieve, states_update] + state_entity_search_tools
return [states_create, states_list, states_retrieve, states_update]
apps/pi/pi/services/actions/tools/types.py
+67 -22
View File
@@ -2,6 +2,8 @@
Types API tools for Plane issue type management.
"""
from typing import Any
from typing import Dict
from typing import Optional
from langchain_core.tools import tool
@@ -20,16 +22,26 @@ def get_type_tools(method_executor, context):
async def types_create(
name: str,
description: Optional[str] = None,
project_ids: Optional[list[str]] = None,
is_epic: Optional[bool] = None,
is_active: Optional[bool] = None,
external_source: Optional[str] = None,
external_id: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Create a new issue type.
) -> Dict[str, Any]:
"""Create a new work item type.
Args:
name: Parameter description (required)
description: Parameter description (optional)
project_id: Parameter description (optional)
workspace_slug: Parameter description (optional)
name: Name of the work item type (required)
description: Optional description of the type
project_ids: List of project IDs to associate with this type
is_epic: Flag to mark this type as an epic type
is_active: Activation status of the type
external_source: External integration source (e.g., 'jira', 'github')
external_id: External system identifier for this type
project_id: UUID of the project (auto-filled from context)
workspace_slug: Workspace slug identifier (auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -42,16 +54,21 @@ def get_type_tools(method_executor, context):
"create",
name=name,
description=description,
project_ids=project_ids,
is_epic=is_epic,
is_active=is_active,
external_source=external_source,
external_id=external_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response(f"Successfully created type '{name}'", result["data"])
return PlaneToolBase.format_success_payload(f"Successfully created type '{name}'", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create type", result["error"])
return PlaneToolBase.format_error_payload("Failed to create type", result["error"])
@tool
async def types_list(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> str:
async def types_list(project_id: Optional[str] = None, workspace_slug: Optional[str] = None) -> Dict[str, Any]:
"""List all work item types."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -61,16 +78,16 @@ def get_type_tools(method_executor, context):
result = await method_executor.execute("types", "list", project_id=project_id, workspace_slug=workspace_slug)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved types list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved types list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list types", result["error"])
return PlaneToolBase.format_error_payload("Failed to list types", result["error"])
@tool
async def types_retrieve(
type_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Get a single type by ID."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -86,19 +103,37 @@ def get_type_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved type", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved type", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve type", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve type", result["error"])
@tool
async def types_update(
type_id: str,
name: Optional[str] = None,
description: Optional[str] = None,
project_ids: Optional[list[str]] = None,
is_epic: Optional[bool] = None,
is_active: Optional[bool] = None,
external_source: Optional[str] = None,
external_id: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Update type details."""
) -> Dict[str, Any]:
"""Update work item type details.
Args:
type_id: UUID of the work item type to update (required)
name: Updated name of the type
description: Updated description
project_ids: Updated list of project IDs to associate with this type
is_epic: Updated epic type flag
is_active: Updated activation status
external_source: Updated external integration source
external_id: Updated external system identifier
project_id: UUID of the project (auto-filled from context)
workspace_slug: Workspace slug identifier (auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -106,11 +141,21 @@ def get_type_tools(method_executor, context):
project_id = context["project_id"]
# Build update data
update_data = {}
update_data: dict[str, Any] = {}
if name is not None:
update_data["name"] = name
if description is not None:
update_data["description"] = description
if project_ids is not None:
update_data["project_ids"] = project_ids
if is_epic is not None:
update_data["is_epic"] = is_epic
if is_active is not None:
update_data["is_active"] = is_active
if external_source is not None:
update_data["external_source"] = external_source
if external_id is not None:
update_data["external_id"] = external_id
result = await method_executor.execute(
"types",
@@ -121,16 +166,16 @@ def get_type_tools(method_executor, context):
**update_data,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated type", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated type", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update type", result["error"])
return PlaneToolBase.format_error_payload("Failed to update type", result["error"])
@tool
async def types_delete(
type_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Delete a type."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -146,8 +191,8 @@ def get_type_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted type", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted type", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete type", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete type", result["error"])
return [types_create, types_list, types_retrieve, types_update, types_delete]
apps/pi/pi/services/actions/tools/users.py
+7 -10
View File
@@ -2,6 +2,9 @@
Users API tools for Plane user management operations.
"""
from typing import Any
from typing import Dict
from langchain_core.tools import tool
from .base import PlaneToolBase
@@ -14,18 +17,12 @@ def get_user_tools(method_executor, context):
"""Return LangChain tools for the users category using method_executor and context."""
@tool
async def users_get_current() -> str:
async def users_get_current() -> Dict[str, Any]:
"""Get current user information."""
result = await method_executor.execute("users", "get_current")
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved current user", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved current user", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to get current user", result["error"])
return PlaneToolBase.format_error_payload("Failed to get current user", result["error"])
# Get entity search tools relevant only to users
from .entity_search import get_entity_search_tools
entity_search_tools = get_entity_search_tools(method_executor, context)
user_entity_search_tools = [t for t in entity_search_tools if getattr(t, "name", "").find("user") != -1]
return [users_get_current] + user_entity_search_tools
return [users_get_current]
apps/pi/pi/services/actions/tools/workitems.py
+190 -97
View File
@@ -4,6 +4,8 @@ Work Items API tools for Plane issue/task management operations.
import logging
import uuid
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
@@ -71,6 +73,51 @@ async def resolve_state_to_uuid(state: Optional[str], project_id: Optional[str],
return None
async def resolve_type_to_uuid(type_id: Optional[str], project_id: Optional[str], workspace_slug: Optional[str] = None) -> Optional[str]:
"""
Resolve issue type name to UUID with enhanced matching strategy.
Args:
type_id: Type name or UUID
project_id: Project ID for type resolution
workspace_slug: Optional workspace slug for filtering
Returns:
Type UUID if resolved, None if type should use project default or if resolution fails
"""
if not type_id:
return None
# Check if type_id is already a valid UUID
try:
uuid.UUID(type_id)
return type_id
except (ValueError, TypeError):
log.debug(f"Type '{type_id}' is not a UUID, will resolve as name")
# Only attempt resolution if we have a project_id
if not project_id:
log.warning(f"Cannot resolve type '{type_id}' without project_id")
return None
try:
from pi.app.api.v1.helpers.plane_sql_queries import search_type_by_name
type_result = await search_type_by_name(type_id, project_id, workspace_slug)
if type_result and "id" in type_result:
resolved_uuid = type_result["id"]
log.info(f"Resolved type '{type_id}' to UUID {resolved_uuid}")
return resolved_uuid
else:
log.warning(f"Could not resolve type '{type_id}' to a valid UUID")
return None
except Exception as e:
log.error(f"Error resolving type '{type_id}': {e}")
return None
async def get_epic_type_id(method_executor, project_id: str, workspace_slug: str) -> Optional[str]:
"""Get the epic issue type ID for a project using direct SQL query.
@@ -117,33 +164,25 @@ def get_workitem_tools(method_executor, context):
target_date: Optional[str] = None,
type_id: Optional[str] = None,
parent: Optional[str] = None,
module_id: Optional[str] = None,
cycle_id: Optional[str] = None,
estimate_point: Optional[str] = None,
point: Optional[int] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
is_draft: Optional[bool] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new work item/issue.
Args:
name: Work item title (required)
project_id: Project ID (required - provide from conversation context or previous actions)
workspace_slug: Workspace slug (required - provide from conversation context)
description_html: Issue description in HTML format
description_html: Workitem description in HTML format
priority: Priority level (high, medium, low, urgent, none)
state: State name or UUID for the issue (e.g., "todo", "in progress", "done")
state: State name or UUID for the workitem (e.g., "todo", "in progress", "done")
assignees: List of assignee user IDs
labels: List of label IDs
start_date: Start date (YYYY-MM-DD format)
target_date: Target completion date (YYYY-MM-DD format)
type_id: Issue type ID (optional)
type_id: Workitem type ID (optional). MUST be the 'issue_types_id' (UUID of the issue type definition), NOT the 'project_issue_types_id'.
parent: Parent work-item ID (optional)
module_id: Module ID to associate (optional)
cycle_id: Cycle ID to associate (optional)
estimate_point: Estimate point ID (optional)
point: Story points value (optional)
external_id: External system identifier (optional)
external_source: External system source name (optional)
is_draft: Create as draft (optional)
@@ -161,6 +200,9 @@ def get_workitem_tools(method_executor, context):
# Resolve state name to UUID if needed
resolved_state = await resolve_state_to_uuid(state, project_id, workspace_slug)
# Resolve type name to UUID if needed
resolved_type_id = await resolve_type_to_uuid(type_id, project_id, workspace_slug)
result = await method_executor.execute(
"workitems",
"create",
@@ -174,12 +216,8 @@ def get_workitem_tools(method_executor, context):
labels=labels,
start_date=start_date,
target_date=target_date,
type_id=type_id,
type_id=resolved_type_id,
parent=parent,
module_id=module_id,
cycle_id=cycle_id,
estimate_point=estimate_point,
point=point,
external_id=external_id,
external_source=external_source,
is_draft=is_draft,
@@ -201,9 +239,14 @@ def get_workitem_tools(method_executor, context):
except Exception as e:
log.warning(f"Could not enrich workitem data with identifier info: {e}")
return await PlaneToolBase.format_success_response_with_url(f"Successfully created work item '{name}'", data, "workitem", context)
return await PlaneToolBase.format_success_payload_with_url(
f"Successfully created work item '{name}'",
data,
"workitem",
context,
)
else:
return PlaneToolBase.format_error_response("Failed to create work item", result["error"])
return PlaneToolBase.format_error_payload("Failed to create work item", result.get("error"))
@tool
async def workitems_update(
@@ -219,14 +262,10 @@ def get_workitem_tools(method_executor, context):
target_date: Optional[str] = None,
type_id: Optional[str] = None,
parent: Optional[str] = None,
module_id: Optional[str] = None,
cycle_id: Optional[str] = None,
estimate_point: Optional[str] = None,
point: Optional[int] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
is_draft: Optional[bool] = None,
) -> str:
) -> Dict[str, Any]:
"""Update an existing work item/issue.
Args:
@@ -241,12 +280,8 @@ def get_workitem_tools(method_executor, context):
labels: New list of label IDs
start_date: New start date (YYYY-MM-DD format)
target_date: New target completion date (YYYY-MM-DD format)
type_id: Issue type ID (optional)
type_id: Workitem type ID (optional). MUST be the 'issue_types_id' (UUID of the issue type definition), NOT the 'project_issue_types_id'.
parent: Parent work-item ID (optional)
module_id: Module ID to associate (optional)
cycle_id: Cycle ID to associate (optional)
estimate_point: Estimate point ID (optional)
point: Story points value (optional)
external_id: External system identifier (optional)
external_source: External system source name (optional)
is_draft: Mark as draft (optional)
@@ -264,6 +299,9 @@ def get_workitem_tools(method_executor, context):
# Resolve state name to UUID if needed
resolved_state = await resolve_state_to_uuid(state, project_id, workspace_slug)
# Resolve type name to UUID if needed
resolved_type_id = await resolve_type_to_uuid(type_id, project_id, workspace_slug)
# Build update data with only non-None values to avoid overwriting existing fields
update_data = {
k: v
@@ -276,12 +314,8 @@ def get_workitem_tools(method_executor, context):
"labels": labels,
"start_date": start_date,
"target_date": target_date,
"type_id": type_id,
"type_id": resolved_type_id,
"parent": parent,
"module_id": module_id,
"cycle_id": cycle_id,
"estimate_point": estimate_point,
"point": point,
"external_id": external_id,
"external_source": external_source,
"is_draft": is_draft,
@@ -314,12 +348,12 @@ def get_workitem_tools(method_executor, context):
except Exception as e:
log.warning(f"Could not enrich workitem update data with identifier info: {e}")
return await PlaneToolBase.format_success_response_with_url("Successfully updated work item", data, "workitem", context)
return await PlaneToolBase.format_success_payload_with_url("Successfully updated work item", data, "workitem", context)
else:
return PlaneToolBase.format_error_response("Failed to update work item", result["error"])
return PlaneToolBase.format_error_payload("Failed to update work item", result["error"])
@tool
async def epics_create(
async def create_epic(
name: str,
project_id: Optional[str] = None,
description_html: Optional[str] = None,
@@ -329,13 +363,9 @@ def get_workitem_tools(method_executor, context):
labels: Optional[List[str]] = None,
start_date: Optional[str] = None,
target_date: Optional[str] = None,
parent: Optional[str] = None,
estimate_point: Optional[str] = None,
point: Optional[int] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
is_draft: Optional[bool] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new epic work item.
Args:
@@ -348,12 +378,8 @@ def get_workitem_tools(method_executor, context):
labels: List of label IDs
start_date: Start date (YYYY-MM-DD format)
target_date: Target completion date (YYYY-MM-DD format)
parent: Parent epic ID (optional)
estimate_point: Estimate point ID (optional)
point: Story points value (optional)
external_id: External system identifier (optional)
external_source: External system source name (optional)
is_draft: Create as draft (optional)
Note: This tool automatically sets the work item type to 'epic' for the specified project.
To add the created epic to a module or cycle, use modules_add_work_items or cycles_add_work_items after creation.
@@ -367,11 +393,11 @@ def get_workitem_tools(method_executor, context):
# Get the epic type ID for this project
if project_id is None:
return PlaneToolBase.format_error_response("Failed to create epic", "Project ID is required for epic creation. Please specify a project.")
return PlaneToolBase.format_error_payload("Failed to create epic", "Project ID is required for epic creation. Please specify a project.")
epic_type_id = await get_epic_type_id(method_executor, project_id, workspace_slug)
if not epic_type_id:
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"Failed to create epic", "Could not find epic issue type for this project. Please ensure an epic issue type exists."
)
@@ -392,12 +418,8 @@ def get_workitem_tools(method_executor, context):
start_date=start_date,
target_date=target_date,
type_id=epic_type_id, # Automatically set to epic type
parent=parent,
estimate_point=estimate_point,
point=point,
external_id=external_id,
external_source=external_source,
is_draft=is_draft,
)
if result["success"]:
@@ -416,12 +438,12 @@ def get_workitem_tools(method_executor, context):
except Exception as e:
log.warning(f"Could not enrich epic data with identifier info: {e}")
return await PlaneToolBase.format_success_response_with_url(f"Successfully created epic '{name}'", data, "epic", context)
return await PlaneToolBase.format_success_payload_with_url(f"Successfully created epic '{name}'", data, "epic", context)
else:
return PlaneToolBase.format_error_response("Failed to create epic", result["error"])
return PlaneToolBase.format_error_payload("Failed to create epic", result["error"])
@tool
async def epics_update(
async def update_epic(
issue_id: str,
project_id: Optional[str] = None,
name: Optional[str] = None,
@@ -432,13 +454,9 @@ def get_workitem_tools(method_executor, context):
labels: Optional[List[str]] = None,
start_date: Optional[str] = None,
target_date: Optional[str] = None,
parent: Optional[str] = None,
estimate_point: Optional[str] = None,
point: Optional[int] = None,
external_id: Optional[str] = None,
external_source: Optional[str] = None,
is_draft: Optional[bool] = None,
) -> str:
) -> Dict[str, Any]:
"""Update an existing epic work item.
Args:
@@ -452,12 +470,8 @@ def get_workitem_tools(method_executor, context):
labels: List of label IDs
start_date: Start date (YYYY-MM-DD format)
target_date: Target completion date (YYYY-MM-DD format)
parent: Parent epic ID (optional)
estimate_point: Estimate point ID (optional)
point: Story points value (optional)
external_id: External system identifier (optional)
external_source: External system source name (optional)
is_draft: Mark as draft (optional)
Note: This tool updates an existing epic. It does not change the work item type.
To add/move the epic to a module or cycle, use modules_add_work_items or cycles_add_work_items.
@@ -485,12 +499,8 @@ def get_workitem_tools(method_executor, context):
"labels": labels,
"start_date": start_date,
"target_date": target_date,
"parent": parent,
"estimate_point": estimate_point,
"point": point,
"external_id": external_id,
"external_source": external_source,
"is_draft": is_draft,
}.items()
if v is not None
}
@@ -505,9 +515,9 @@ def get_workitem_tools(method_executor, context):
)
if result["success"]:
return await PlaneToolBase.format_success_response_with_url("Successfully updated epic", result["data"], "epic", context)
return await PlaneToolBase.format_success_payload_with_url("Successfully updated epic", result["data"], "epic", context)
else:
return PlaneToolBase.format_error_response("Failed to update epic", result["error"])
return PlaneToolBase.format_error_payload("Failed to update epic", result["error"])
@tool
async def workitems_list(
@@ -520,7 +530,22 @@ def get_workitem_tools(method_executor, context):
fields: Optional[str] = None,
order_by: Optional[str] = None,
per_page: Optional[int] = 20,
) -> str:
priority: Optional[str] = None,
state: Optional[str] = None,
assignees: Optional[List[str]] = None,
labels: Optional[List[str]] = None,
start_date: Optional[str] = None,
target_date: Optional[str] = None,
created_by: Optional[str] = None,
updated_by: Optional[str] = None,
type_id: Optional[str] = None,
parent: Optional[str] = None,
is_draft: Optional[bool] = None,
created_at: Optional[str] = None,
updated_at: Optional[str] = None,
completed_at: Optional[str] = None,
archived_at: Optional[str] = None,
) -> Dict[str, Any]:
"""List work items with filtering.
Args:
@@ -533,6 +558,21 @@ def get_workitem_tools(method_executor, context):
fields: Comma-separated list of fields to include in response
order_by: Field to order results by. Prefix with '-' for descending order
per_page: Number of work items per page (default: 20, max: 100)
priority: Filter by priority (high, medium, low, urgent, none)
state: Filter by state name or UUID
assignees: Filter by assignee user IDs
labels: Filter by label IDs
start_date: Filter by start date
target_date: Filter by target date
created_by: Filter by creator ID
updated_by: Filter by updater ID
type_id: Filter by issue type ID
parent: Filter by parent issue ID
is_draft: Filter by draft status
created_at: Filter by creation time
updated_at: Filter by update time
completed_at: Filter by completion time
archived_at: Filter by archive time
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -540,6 +580,31 @@ def get_workitem_tools(method_executor, context):
if project_id is None and "project_id" in context:
project_id = context["project_id"]
# Resolve state name to UUID if needed
resolved_state = await resolve_state_to_uuid(state, project_id, workspace_slug)
# Build filter arguments
filter_args = {
"priority": priority,
"state": resolved_state,
"assignees": assignees,
"labels": labels,
"start_date": start_date,
"target_date": target_date,
"created_by": created_by,
"updated_by": updated_by,
"type_id": type_id,
"parent": parent,
"is_draft": is_draft,
"created_at": created_at,
"updated_at": updated_at,
"completed_at": completed_at,
"archived_at": archived_at,
}
# Remove None values
filter_args = {k: v for k, v in filter_args.items() if v is not None}
result = await method_executor.execute(
"workitems",
"list",
@@ -555,13 +620,13 @@ def get_workitem_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved work items list", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved work items list", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list work items", result["error"])
return PlaneToolBase.format_error_payload("Failed to list work items", result["error"])
@tool
async def workitems_retrieve(
pk: str,
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
expand: Optional[str] = None,
@@ -569,11 +634,11 @@ def get_workitem_tools(method_executor, context):
external_source: Optional[str] = None,
fields: Optional[str] = None,
order_by: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Retrieve a single work item by ID.
Args:
pk: Work item ID (required)
issue_id: Work item ID (required)
project_id: Project ID (required - provide from conversation context or previous actions)
workspace_slug: Workspace slug (provide if known, otherwise auto-detected)
expand: Comma-separated list of related fields to expand in response
@@ -591,7 +656,7 @@ def get_workitem_tools(method_executor, context):
result = await method_executor.execute(
"workitems",
"retrieve",
pk=pk,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
expand=expand,
@@ -602,9 +667,9 @@ def get_workitem_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved work item", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved work item", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve work item", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve work item", result["error"])
@tool
async def workitems_search(
@@ -613,7 +678,7 @@ def get_workitem_tools(method_executor, context):
limit: Optional[int] = None,
project_id: Optional[str] = None,
workspace_search: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Search work items by criteria.
Args:
@@ -640,16 +705,16 @@ def get_workitem_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully searched work items", result["data"])
return PlaneToolBase.format_success_payload("Successfully searched work items", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to search work items", result["error"])
return PlaneToolBase.format_error_payload("Failed to search work items", result["error"])
@tool
async def workitems_get_workspace(
issue_identifier: int,
project_identifier: str,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Get work item across workspace using identifiers.
Args:
@@ -670,9 +735,9 @@ def get_workitem_tools(method_executor, context):
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved workspace work item", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved workspace work item", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to retrieve workspace work item", result["error"])
return PlaneToolBase.format_error_payload("Failed to retrieve workspace work item", result["error"])
@tool
async def workitems_create_relation(
@@ -680,7 +745,7 @@ def get_workitem_tools(method_executor, context):
relation_type: str,
related_issues: List[str],
project_id: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Create relationships between work items.
Args:
@@ -714,13 +779,13 @@ def get_workitem_tools(method_executor, context):
# Validate relation type
if relation_type not in RELATION_TYPES:
valid_types = ", ".join(RELATION_TYPES.keys())
return PlaneToolBase.format_error_response(
return PlaneToolBase.format_error_payload(
"Invalid relation type", f"Relation type '{relation_type}' is not valid. Valid types are: {valid_types}"
)
# Validate required parameters
if not related_issues:
return PlaneToolBase.format_error_response("Missing related issues", "At least one related issue ID must be provided")
return PlaneToolBase.format_error_payload("Missing related issues", "At least one related issue ID must be provided")
result = await method_executor.execute(
"workitems",
@@ -739,26 +804,54 @@ def get_workitem_tools(method_executor, context):
# Create a mock work item response with the source issue_id for entity URL generation
source_workitem_data = {"id": issue_id}
return await PlaneToolBase.format_success_response_with_url(
return await PlaneToolBase.format_success_payload_with_url(
f"Successfully created {relation_count} '{relation_type}' relation(s) for work item", source_workitem_data, "workitem", context
)
else:
return PlaneToolBase.format_error_response("Failed to create work item relation", result["error"])
return PlaneToolBase.format_error_payload("Failed to create work item relation", result["error"])
# Get entity search tools relevant only to workitems
from .entity_search import get_entity_search_tools
@tool
async def workitems_delete(
issue_id: str,
project_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Delete a work item/issue.
entity_search_tools = get_entity_search_tools(method_executor, context)
workitem_entity_search_tools = [t for t in entity_search_tools if "workitem" in t.name or "user" in t.name]
Args:
issue_id: Work item ID to delete (required)
project_id: Project ID (required - provide from conversation context or previous actions)
Warning: This action is permanent and cannot be undone. The work item will be permanently deleted.
"""
# Auto-fill workspace_slug from context (hidden from LLM)
workspace_slug = context.get("workspace_slug")
# Auto-fill project_id from context if not provided
if project_id is None and "project_id" in context:
project_id = context["project_id"]
result = await method_executor.execute(
"workitems",
"delete",
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_payload(f"Successfully deleted work item {issue_id}", {"id": issue_id})
else:
return PlaneToolBase.format_error_payload("Failed to delete work item", result.get("error"))
return [
workitems_create,
workitems_update,
epics_create,
epics_update,
create_epic,
update_epic,
workitems_create_relation,
workitems_delete,
workitems_list,
# workitems_retrieve,
# workitems_search,
# workitems_get_workspace
] + workitem_entity_search_tools
]
apps/pi/pi/services/actions/tools/worklogs.py
+46 -21
View File
@@ -22,17 +22,21 @@ def get_worklog_tools(method_executor, context):
issue_id: str,
description: str,
duration: int,
created_by: Optional[str] = None,
updated_by: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Create a new worklog entry.
Args:
issue_id: Parameter description (required)
description: Parameter description (required)
duration: Parameter description (required)
project_id: Parameter description (optional)
workspace_slug: Parameter description (optional)
issue_id: Issue ID (required)
description: Worklog description (required)
duration: Duration in minutes (required)
created_by: User ID who created the worklog
updated_by: User ID who updated the worklog
project_id: Project ID (optional, auto-filled from context)
workspace_slug: Workspace slug (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -46,20 +50,22 @@ def get_worklog_tools(method_executor, context):
issue_id=issue_id,
description=description,
duration=duration,
created_by=created_by,
updated_by=updated_by,
project_id=project_id,
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully created worklog", result["data"])
return PlaneToolBase.format_success_payload("Successfully created worklog", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to create worklog", result["error"])
return PlaneToolBase.format_error_payload("Failed to create worklog", result["error"])
@tool
async def worklogs_list(
issue_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""List worklogs for an issue."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -75,15 +81,15 @@ def get_worklog_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved worklogs", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved worklogs", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to list worklogs", result["error"])
return PlaneToolBase.format_error_payload("Failed to list worklogs", result["error"])
@tool
async def worklogs_get_summary(
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Get project worklog summary."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -98,19 +104,33 @@ def get_worklog_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully retrieved worklog summary", result["data"])
return PlaneToolBase.format_success_payload("Successfully retrieved worklog summary", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to get worklog summary", result["error"])
return PlaneToolBase.format_error_payload("Failed to get worklog summary", result["error"])
@tool
async def worklogs_update(
worklog_id: str,
issue_id: str,
description: Optional[str] = None,
duration: Optional[int] = None,
created_by: Optional[str] = None,
updated_by: Optional[str] = None,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
"""Update time entry."""
) -> Dict[str, Any]:
"""Update time entry.
Args:
worklog_id: Worklog ID (required)
issue_id: Issue ID (required)
description: Worklog description
duration: Duration in minutes
created_by: User ID who created the worklog
updated_by: User ID who updated the worklog
project_id: Project ID (optional, auto-filled from context)
workspace_slug: Workspace slug (optional, auto-filled from context)
"""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
workspace_slug = context["workspace_slug"]
@@ -123,26 +143,31 @@ def get_worklog_tools(method_executor, context):
update_data["description"] = description
if duration is not None:
update_data["duration"] = duration
if created_by is not None:
update_data["created_by"] = created_by
if updated_by is not None:
update_data["updated_by"] = updated_by
result = await method_executor.execute(
"worklogs",
"update",
worklog_id=worklog_id,
issue_id=issue_id,
project_id=project_id,
workspace_slug=workspace_slug,
**update_data,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully updated worklog", result["data"])
return PlaneToolBase.format_success_payload("Successfully updated worklog", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to update worklog", result["error"])
return PlaneToolBase.format_error_payload("Failed to update worklog", result["error"])
@tool
async def worklogs_delete(
worklog_id: str,
project_id: Optional[str] = None,
workspace_slug: Optional[str] = None,
) -> str:
) -> Dict[str, Any]:
"""Delete time entry."""
# Auto-fill from context if not provided
if workspace_slug is None and "workspace_slug" in context:
@@ -158,8 +183,8 @@ def get_worklog_tools(method_executor, context):
workspace_slug=workspace_slug,
)
if result["success"]:
return PlaneToolBase.format_success_response("Successfully deleted worklog", result["data"])
return PlaneToolBase.format_success_payload("Successfully deleted worklog", result["data"])
else:
return PlaneToolBase.format_error_response("Failed to delete worklog", result["error"])
return PlaneToolBase.format_error_payload("Failed to delete worklog", result["error"])
return [worklogs_create, worklogs_list, worklogs_get_summary, worklogs_update, worklogs_delete]
apps/pi/pi/services/chat/action_executor.py
+199
View File
@@ -0,0 +1,199 @@
# batch_execution.py (~150 lines total)
import asyncio
from datetime import datetime
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi import settings
from pi.app.api.v1.helpers.plane_sql_queries import get_workspace_slug
from pi.services.actions import MethodExecutor
from pi.services.actions import PlaneActionsExecutor
from pi.services.chat.chat import PlaneChatBot
from pi.services.chat.helpers.action_execution_helpers import format_response
from pi.services.chat.helpers.action_execution_helpers import load_artifacts
from pi.services.chat.helpers.action_execution_helpers import update_flow_steps
from pi.services.chat.helpers.placeholder_orchestrator import PlaceholderOrchestrator
log = logger.getChild(__name__)
async def get_execution_context(request, chatbot: PlaneChatBot, user_id, db) -> Tuple[str, Optional[str], bool]:
"""Extract execution context from request and database."""
# Token
if request.access_token:
token = request.access_token
log.info(f"Using access_token from request for user {user_id}")
else:
token = await chatbot._get_oauth_token_for_user(db, str(user_id), str(request.workspace_id))
if token:
log.info(f"Retrieved OAuth token from database for user {user_id}, workspace {request.workspace_id}")
else:
log.warning(f"No valid OAuth token found for user {user_id}, workspace {request.workspace_id}")
# Project ID from preferences
from pi.services.retrievers.pg_store.chat import get_project_from_chat_preference
project_id, is_project_chat = await get_project_from_chat_preference(request.chat_id, user_id, db)
return token, project_id, is_project_chat
class BuildModeToolExecutor:
def __init__(self, chatbot: PlaneChatBot, db: AsyncSession):
self.chatbot = chatbot
self.db = db
async def execute(self, request, user_id):
"""Execute the planned actions"""
log.info(f"EXECUTE ACTION REQUEST: {request}")
# 1. Load artifacts
message_id = request.message_id
chat_id = request.chat_id
planned_actions, original_query, conversation_context = await load_artifacts(request.artifact_data, self.db)
# 2. Get context
workspace_id = request.workspace_id
workspace_slug = planned_actions[0]["args"].get("workspace_slug", "")
if not workspace_slug:
workspace_slug = await get_workspace_slug(str(workspace_id))
token, project_id, _ = await get_execution_context(request, self.chatbot, user_id, self.db)
# Validate token exists
if not token:
error_msg = f"No valid OAuth token found for user {user_id} and workspace {request.workspace_id}. Please re-authorize Plane AI."
log.error(error_msg)
raise ValueError(error_msg)
# Log token type for debugging (without exposing the actual token)
token_type = "API key" if token.startswith("plane_api_") else "OAuth access token"
log.info(f"Using {token_type} for execution (token length: {len(token)}, starts with: {token[:10]}...)")
# 3. Execution
# Setup
# Create actions executor
if token_type == "API key":
actions_executor = PlaneActionsExecutor(api_key=token, base_url=settings.plane_api.HOST)
else:
actions_executor = PlaneActionsExecutor(access_token=token, base_url=settings.plane_api.HOST)
method_executor = MethodExecutor(actions_executor)
context = {"workspace_slug": workspace_slug, "project_id": project_id, "message_id": message_id, "chat_id": chat_id}
# Execute single action
# If single action, execute directly
if len(planned_actions) == 1:
log.info(f"Single action execution: {planned_actions[0]}")
results = await self.execute_single_action(planned_actions[0], method_executor, context)
elif self._are_actions_independent(planned_actions):
# Execute independent actions in parallel
log.info(f"Parallel execution for {len(planned_actions)} independent actions")
tasks = [self.execute_single_action(action, method_executor, context) for action in planned_actions]
# Gather results and flatten the list of lists
parallel_results = await asyncio.gather(*tasks)
results = [item for sublist in parallel_results for item in sublist]
else:
# Orchestrate tool execution
# If more than one action, orchestrate tool execution
results = await self.orchestrate_execution(original_query, planned_actions, conversation_context, method_executor, context)
# 4. Update DB and return
await update_flow_steps(results, message_id, chat_id, self.db)
return format_response(planned_actions, results, datetime.utcnow())
async def execute_single_action(self, planned_action, method_executor, context):
"""Execute a single action."""
tool_name = planned_action["tool_name"]
args = planned_action["args"]
entity_type = planned_action["entity_type"]
# Build tool
all_category_tools = self.chatbot._build_method_tools(entity_type, method_executor, context)
# get the specific tool we need
tool = next((t for t in all_category_tools if t.name == tool_name), None)
if not tool:
raise ValueError(f"Tool '{tool_name}' not found in category '{entity_type}'")
# Execute the tool - EXPECT structured payload dict
result = await tool.ainvoke(args)
log.info(f"Tool '{tool_name}' executed with payload: {result}")
if not isinstance(result, dict):
raise ValueError(f"Tool '{tool_name}' must return a structured payload dict, got {type(result)}")
# Normalize dict payload
message = result.get("message") or ""
ok = bool(result.get("ok", True))
entity_info = result.get("entity")
executed = {
"tool_name": tool_name,
"result": message,
"entity_info": entity_info,
"artifact_id": planned_action.get("artifact_id"),
"sequence": 1,
"artifact_type": entity_type,
"executed_at": datetime.utcnow().isoformat(),
"success": ok,
}
return [executed]
async def orchestrate_execution(self, original_query, planned_actions, conversation_context, method_executor, context):
"""
Execute actions using hybrid orchestration with dependency resolution.
This method uses PlaceholderOrchestrator to:
- Programmatically track dependencies between actions
- Execute independent actions when ready
- Use focused LLM calls to extract specific values from tool results
- Validate extracted values (e.g., UUID format)
"""
log.info(f"\n\nORCHESTRATE EXECUTION: {len(planned_actions)} actions\n")
# Use the new hybrid orchestrator
orchestrator = PlaceholderOrchestrator(
planned_actions=planned_actions,
chatbot=self.chatbot,
method_executor=method_executor,
context=context,
chat_id=str(context.get("chat_id")),
message_id=str(context.get("message_id")),
db=self.db,
)
# Execute all actions with dependency resolution
results = await orchestrator.execute_all()
log.info(f"Orchestration complete: {len(results)} actions executed")
return results
def _has_placeholder(self, value: Any) -> bool:
"""Check if a value contains a placeholder recursively."""
if isinstance(value, str):
return "<id of" in value
elif isinstance(value, dict):
return any(self._has_placeholder(v) for v in value.values())
elif isinstance(value, list):
return any(self._has_placeholder(v) for v in value)
return False
def _are_actions_independent(self, planned_actions: List[Dict]) -> bool:
"""Check if all planned actions are independent (no placeholders)."""
for action in planned_actions:
if self._has_placeholder(action.get("args", {})):
return False
return True
apps/pi/pi/services/chat/action_planner.py
+872
View File
@@ -0,0 +1,872 @@
"""Action execution logic with retrieval tools."""
import contextlib
import json
from collections.abc import AsyncIterator
from typing import Any
from typing import Dict
from typing import Optional
from typing import Sequence
from typing import Union
from typing import cast
from langchain_core.messages import HumanMessage
from langchain_core.messages import SystemMessage
from pi import logger
from pi import settings
from pi.app.models.enums import MessageMetaStepType
from pi.core.db.plane_pi.lifecycle import get_streaming_db_session
from pi.services.chat.utils import get_current_timestamp_context
from pi.services.chat.utils import reasoning_dict_maker
from pi.services.retrievers.pg_store.message import upsert_message_flow_steps as _upsert_message_flow_steps
from pi.services.schemas.chat import ActionCategorySelection
from .helpers.build_mode_helpers import TOOL_NAME_TO_CATEGORY_MAP
from .helpers.build_mode_helpers import auto_resolve_missing_ids
from .helpers.build_mode_helpers import build_advisory_tool_step
from .helpers.build_mode_helpers import build_and_stream_final_response
from .helpers.build_mode_helpers import build_method_executor_and_context
from .helpers.build_mode_helpers import build_planner_summary_step
from .helpers.build_mode_helpers import build_planner_tool_selection_step
from .helpers.build_mode_helpers import build_planning_tools
from .helpers.build_mode_helpers import build_tool_orchestration_context_step
from .helpers.build_mode_helpers import enrich_tool_query_for_display
from .helpers.build_mode_helpers import execute_and_persist_clarification
from .helpers.build_mode_helpers import execute_retrieval_tool_and_build_step
from .helpers.build_mode_helpers import handle_preflight_clarification
from .helpers.build_mode_helpers import persist_skip_category_selection_step
from .helpers.build_mode_helpers import plan_action_and_prepare_outputs
from .helpers.build_mode_helpers import recover_clarification_categories
from .helpers.build_mode_helpers import run_category_router_and_persist
from .helpers.build_mode_helpers import selected_action_categories_display
from .helpers.tool_utils import build_method_prompt
from .helpers.tool_utils import classify_tool
from .helpers.tool_utils import format_tool_message_for_display
from .helpers.tool_utils import format_tool_query_for_display
from .helpers.tool_utils import preflight_missing_required_fields
from .helpers.tool_utils import stream_content_in_chunks
from .helpers.tool_utils import tool_name_shown_to_user
# from .tool_utils import log_toolset_details
log = logger.getChild(__name__)
MAX_ACTION_EXECUTOR_ITERATIONS = settings.chat.MAX_ACTION_EXECUTOR_ITERATIONS
async def execute_tools_for_build_mode(
chatbot_instance,
selected_tools,
user_meta,
workspace_id,
workspace_slug,
project_id,
conversation_history,
enhanced_conversation_history, # 🆕 Enhanced context with action details
user_id,
chat_id,
query_flow_store,
combined_tool_query,
query_id,
step_order,
db,
reasoning_container=None,
is_project_chat=None,
pi_sidebar_open=None,
sidebar_open_url=None,
) -> AsyncIterator[Union[str, Dict[str, Any]]]:
"""
Execute action planner with access to retrieval tools
Steps:
1. Category Selection (Programmatic + LLM Router)
2. Method Planning & Execution
3. Tool Execution
4. Final Response
"""
try:
# Resolve workspace_id from project_id if needed (for project-level chats)
# This must happen BEFORE Phase 1 tools are built
if not workspace_id and project_id:
try:
from pi.app.api.v1.helpers.plane_sql_queries import resolve_workspace_id_from_project_id
resolved_workspace_id = await resolve_workspace_id_from_project_id(str(project_id))
if resolved_workspace_id:
workspace_id = str(resolved_workspace_id)
log.info(f"ChatID: {chat_id} - Resolved workspace_id {workspace_id} from project_id {project_id} (Phase 1)")
except Exception as e:
log.error(f"ChatID: {chat_id} - Failed to resolve workspace_id from project_id {project_id}: {e}")
clarification_requested = False
clarification_payload: dict | None = None
# Build sdk method caller (executor) early for both Phase 1 (advisory) and Phase 2 (planning)
method_executor, context, workspace_slug = await build_method_executor_and_context(
chatbot_instance=chatbot_instance,
user_id=user_id,
workspace_id=workspace_id,
project_id=project_id,
conversation_history=conversation_history,
user_meta=user_meta,
is_project_chat=is_project_chat,
chat_id=chat_id,
db=db,
)
# Update project_id from normalized context
project_id = context["project_id"]
current_step = step_order
tool_flow_steps = []
# Check if resuming after clarification and recover categories
selections_list: Sequence[Union[ActionCategorySelection, Dict[str, Optional[str]]]]
skip_category_selection, selections_list, clar_ctx = await recover_clarification_categories(
user_meta=user_meta,
chat_id=chat_id,
db=db,
)
if not skip_category_selection:
# ----- PHASE 1: Category Selection (Programmatic + LLM Router) -----
# Yield early reasoning chunk to provide immediate feedback (similar to ask mode)
stage = "build_beginning"
content = ""
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content=content)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
## change point: "πspecial reasoning blockπ:
yield reasoning_chunk_dict
# Advisory step (direct function call, not tool invocation)
try:
advisory_text, advisory_step, current_step = await build_advisory_tool_step(
combined_tool_query=combined_tool_query,
current_step=current_step,
)
tool_flow_steps.append(advisory_step)
except Exception as e:
log.warning(f"Category advisory generation failed: {str(e)}")
msg = "An unexpected error occurred. Please try again later."
yield msg
yield f"__FINAL_RESPONSE__{msg}"
return
# Router and persistence (LLM call)
selections_list, current_step = await run_category_router_and_persist(
chatbot_instance=chatbot_instance,
advisory_text=advisory_text,
combined_tool_query=combined_tool_query,
enhanced_conversation_history=enhanced_conversation_history,
query_id=query_id,
chat_id=chat_id,
current_step=current_step,
db=db,
)
# Stream category routing decisions as reasoning for frontend visibility
try:
if selections_list:
try:
content = selected_action_categories_display(selections_list)
stage = "selected_action_categories"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content=content)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
## change point: "πspecial reasoning blockπ:
yield reasoning_chunk_dict
except Exception as e:
log.info(f"ChatID: {chat_id} - Error in reasoning chunk creation in selected_action_categories: {e}", exc_info=True)
pass
except Exception as e:
log.info(f"ChatID: {chat_id} - Error in advisory step in selected_action_categories: {e}", exc_info=True)
pass
if not selections_list:
log.info(f"ChatID: {chat_id} - Category router returned empty list - unsupported action request")
else:
# Record that we skipped category selection due to clarification context
# Yield early reasoning chunk for clarification follow-ups
stage = "actions_clarification_followup"
content = ""
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content=content)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
## change point: "πspecial reasoning blockπ:
yield reasoning_chunk_dict
current_step = await persist_skip_category_selection_step(
selections_list=cast(Sequence[Dict[str, Optional[str]]], selections_list),
current_step=current_step,
enhanced_conversation_history=enhanced_conversation_history,
query_id=query_id,
chat_id=chat_id,
db=db,
)
# ----- PHASE 2: Method Planning & Execution -----
# method_executor already built above; now build the tools
fresh_retrieval_tools = chatbot_instance._create_tools(
db,
user_meta,
workspace_id,
project_id,
user_id,
chat_id,
query_flow_store,
conversation_history,
query_id,
is_project_chat=is_project_chat,
)
combined_tools, all_method_tools, built_categories = build_planning_tools(
chatbot_instance=chatbot_instance,
selections_list=selections_list,
method_executor=method_executor,
context=context,
fresh_retrieval_tools=fresh_retrieval_tools,
)
if not combined_tools:
log.warning("No method or retrieval tools available for selected categories")
msg = "An unexpected error occurred. Please try again later."
yield msg
yield f"__FINAL_RESPONSE__{msg}"
return
# Build the planning prompt via shared helper, pass previously derived clar_ctx
method_prompt = build_method_prompt(
combined_tool_query,
project_id,
user_id,
workspace_id,
enhanced_conversation_history,
clarification_context=clar_ctx,
)
date_time_context = await get_current_timestamp_context(user_id)
method_prompt = f"{method_prompt}\n\n{date_time_context}"
# Record the tool orchestration context (enhanced conversation history) before planning
try:
step_ctx, next_step_ctx = build_tool_orchestration_context_step(
current_step=current_step,
enhanced_conversation_history=enhanced_conversation_history,
combined_tool_query=combined_tool_query,
built_categories=built_categories,
all_method_tools=all_method_tools,
combined_tools=combined_tools,
)
if step_ctx is not None:
tool_flow_steps.append(step_ctx)
current_step = next_step_ctx
except Exception:
pass
# Log the full method planning prompt and context for debugging
try:
# Determine if this is a clarification follow-up (raw user input) or router-synthesized query
is_clarification_followup = bool(user_meta and isinstance(user_meta, dict) and user_meta.get("clarification_context"))
query_label = "User Intent (clarification response)" if is_clarification_followup else "User Intent"
# Log comprehensive debugging information
log.info(f"ChatID: {chat_id} - {query_label}: {combined_tool_query}")
log.info(f"ChatID: {chat_id} - Selected Categories: {built_categories}")
log.info(f"ChatID: {chat_id} - Available Tools Count: {len(combined_tools)}")
try:
tool_names_for_log = [t.name for t in combined_tools]
log.info(f"ChatID: {chat_id} - Planning tools: {tool_names_for_log}")
except Exception:
pass
except Exception as e:
log.warning(f"ChatID: {chat_id} - Failed to log debug info: {e}")
# Initialize messages for Phase 2 tool orchestration
messages = [SystemMessage(content=method_prompt), HumanMessage(content=combined_tool_query)]
# Re-bind LLM with the full toolset (action methods + retrieval)
# Some LangChain/OpenAI versions default to no tool calls if not specified.
# deduplicate tools
combined_tools = list({tool.name: tool for tool in combined_tools}.values())
llm_with_method_tools = chatbot_instance.tool_llm.bind_tools(combined_tools)
# Set tracking context for method planning LLM calls
llm_with_method_tools.set_tracking_context(query_id, db, MessageMetaStepType.ACTION_METHOD_PLANNING, chat_id=str(chat_id))
stage = "planner_tool_selection_calling"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
# Continue iterative tool calling with method tools
response = await llm_with_method_tools.ainvoke(messages)
messages.append(response)
# Stream planner's initial reasoning (if present and not a control token)
# Only stream as reasoning if there are tool calls coming (not the final response)
# This represents the LLM's internal thinking/reasoning process
# Log LLM's initial response for debugging
log.info(f"ChatID: {chat_id} - LLM INITIAL RESPONSE")
try:
_has_tool_calls = hasattr(response, "tool_calls") and bool(getattr(response, "tool_calls", None))
except Exception:
_has_tool_calls = False
log.info(f"ChatID: {chat_id} - Has Tool Calls: {_has_tool_calls}")
if _has_tool_calls:
# Yield reasoning chunk for the planner tool selection
try:
_content_preview = str(getattr(response, "content", "") or "").strip()
if _content_preview:
## change point: "πspecial reasoning blockπ:
stage = "planner_tool_selection"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content=_content_preview)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
except Exception:
pass
# Log planner decisions immediately (before flow step creation)
try:
if _has_tool_calls:
tool_calls_list = getattr(response, "tool_calls", []) or []
tool_names = []
for _tc in tool_calls_list:
if isinstance(_tc, dict):
tool_names.append(_tc.get("name", ""))
else:
tool_names.append(getattr(_tc, "name", ""))
reasoning_text = str(getattr(response, "content", "") or "").strip()
reason_preview = reasoning_text or "(none)"
log.info(f"ChatID: {chat_id} - Planner selected tools: {tool_names}")
log.info(f"ChatID: {chat_id} - Planner reasoning: {reason_preview}")
else:
# Log when no tool calls are made
reasoning_text = str(getattr(response, "content", "") or "").strip()
log.info(f"ChatID: {chat_id} - Planner selected tools: []")
log.info(f"ChatID: {chat_id} - Planner reasoning (no tools): {reasoning_text}")
except Exception as e:
log.warning(f"ChatID: {chat_id} - Failed to log planner decisions: {e}")
# Record planner decision for this response (selected tools + reasoning)
try:
selection_step, next_step_sel = build_planner_tool_selection_step(response, current_step=current_step)
if selection_step is not None:
tool_flow_steps.append(selection_step)
current_step = next_step_sel
except Exception:
pass
# Handle iterative method tool calling until completion
planned_actions: list = []
max_iterations = MAX_ACTION_EXECUTOR_ITERATIONS # Safety limit to prevent infinite loops
iteration_count = 0
loop_warning_detected = False
# Continue the loop until either:
# 1. No more tool calls are returned by the LLM, OR
# 2. At least one action has been planned (to handle simple requests)
# 3. Maximum iterations reached (safety check)
loop_condition = (_has_tool_calls or len(planned_actions) == 0) and iteration_count < max_iterations
# Track tool calls to detect loops
tool_call_history = []
while loop_condition:
iteration_count += 1
tool_messages = []
# If no tool calls in this response, the LLM has finished its work
# Either it has planned actions, or it's providing a final answer (retrieval-only)
if not _has_tool_calls:
log.info(f"ChatID: {chat_id} - No tool calls returned; LLM has finished (planned {len(planned_actions)} actions)")
break
# Execute all tool calls in this round
log.info(f"ChatID: {chat_id} - Entering the for loop to execute selected tool calls in this iteration: {iteration_count}")
for tool_call in getattr(response, "tool_calls", []):
# Handle both dictionary and object access patterns for LangChain tool calls
if isinstance(tool_call, dict):
tool_name = tool_call.get("name", "")
tool_args = tool_call.get("args", {})
tool_id = tool_call.get("id", "")
else:
# Handle object access pattern
tool_name = getattr(tool_call, "name", "")
tool_args = getattr(tool_call, "args", {})
tool_id = getattr(tool_call, "id", "")
# Additional validation: ensure we have a valid tool name
if not tool_name or not isinstance(tool_name, str):
log.warning(f"Invalid tool name: {tool_name}, skipping tool call")
continue
# Log each tool call execution for debugging
log.info(f"ChatID: {chat_id} - EXECUTING TOOL: {tool_name} with args: {tool_args}")
# Track tool calls to detect loops
tool_call_signature = f"{tool_name}({tool_args})"
tool_call_history.append(tool_call_signature)
# Check for repeated identical calls (potential loop)
if tool_call_history.count(tool_call_signature) > 2:
log.warning(
f"ChatID: {chat_id} - DETECTED LOOP: Tool {tool_name} called {tool_call_history.count(tool_call_signature)} times with same args: {tool_args}" # noqa: E501
) # noqa: E501
log.warning(f"ChatID: {chat_id} - Tool call history: {tool_call_history[-10:]}") # Last 10 calls
loop_warning_detected = True
# Special handling for clarification tool
elif tool_name == "ask_for_clarification":
flow_step, tool_message, next_step, stream_chunk = await execute_and_persist_clarification(
tool_args=tool_args,
tool_id=tool_id,
combined_tools=combined_tools,
combined_tool_query=combined_tool_query,
current_step=current_step,
built_categories=list(built_categories) if "built_categories" in locals() else [],
all_method_tools=list(all_method_tools) if "all_method_tools" in locals() else [],
chat_id=chat_id,
query_id=query_id,
db=db,
)
if tool_message is not None:
tool_messages.append(tool_message)
if flow_step is not None:
tool_flow_steps.append(flow_step)
current_step = next_step
yield stream_chunk
clarification_requested = True
response.tool_calls = [] if hasattr(response, "tool_calls") else None
break
# Check if this is an action tool (not a retrieval tool)
#
# Tool Classification Logic:
# 1. Standard retrieval tools: vector_search_tool, structured_db_tool, etc.
# 2. Read-only operations: tools with patterns like _list, _retrieve, _get, _search
# 3. Modifying operations: tools with patterns like _create, _update, _delete, _add, _remove, _archive, _unarchive
# 4. Safety rule: if a tool has both read-only and modifying patterns, treat as modifying (safer)
#
# This ensures that operations like 'list_modules' don't require user approval,
# while operations like 'create_workitem' still do.
# Classify tool via shared helper
_is_retrieval_tool, is_action_tool = classify_tool(tool_name)
if is_action_tool:
# Build context for entity resolution
action_context = {
"workspace_slug": workspace_slug,
"project_id": project_id,
"workspace_id": workspace_id,
}
# Preflight: ensure required fields exist before planning actions
try:
missing_required = preflight_missing_required_fields(tool_name, tool_args, action_context)
except Exception:
missing_required = []
# Attempt auto-resolution before asking user
did_resolve = await auto_resolve_missing_ids(
tool_name=tool_name,
tool_args=tool_args,
missing_required=missing_required,
combined_tools=combined_tools,
workspace_slug=workspace_slug,
)
if did_resolve:
with contextlib.suppress(Exception):
missing_required = preflight_missing_required_fields(tool_name, tool_args, action_context)
if missing_required:
clar_req, clar_payload, clar_tool_msg, clar_flow_step, next_step = await handle_preflight_clarification(
tool_name=tool_name,
tool_args=tool_args,
action_context=action_context,
missing_required=missing_required,
method_executor=method_executor,
workspace_slug=workspace_slug,
chat_id=chat_id,
tool_id=tool_id,
current_step=current_step,
combined_tool_query=combined_tool_query,
is_project_chat=is_project_chat,
built_categories=list(built_categories) if "built_categories" in locals() else [],
all_method_tools=list(all_method_tools) if "all_method_tools" in locals() else [],
combined_tools=combined_tools,
query_id=query_id,
db=db,
)
if clar_req:
if clar_tool_msg is not None:
tool_messages.append(clar_tool_msg)
if clar_flow_step is not None:
tool_flow_steps.append(clar_flow_step)
current_step = next_step
clarification_requested = True
try:
yield f"πspecial clarification blockπ: {json.dumps(clar_payload)}\n"
except Exception:
yield f"πspecial clarification blockπ: {str(clar_payload)}\n"
response.tool_calls = [] if hasattr(response, "tool_calls") else None
break
stage = "planner_tool_selection_final"
user_friendly_tool_name = TOOL_NAME_TO_CATEGORY_MAP.get(tool_name, {}).get("front_facing_name", tool_name)
if "name" in tool_args or "title" in tool_args:
tq = f"{tool_args.get("name", tool_args.get("title", ""))}"
else:
tq = ""
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name=user_friendly_tool_name, tool_query=tq, content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
planned_step, action_summary, tool_message_ack, next_step = await plan_action_and_prepare_outputs(
tool_name=tool_name,
tool_args=tool_args,
tool_id=tool_id,
current_step=current_step,
chat_id=chat_id,
query_id=query_id,
conversation_history=conversation_history,
combined_tool_query=combined_tool_query,
db=db,
project_id=project_id,
workspace_slug=workspace_slug,
)
# log planned step
yield f"πspecial actions blockπ: {json.dumps(action_summary, default=str)}"
# Record the planned action and tool execution result
if tool_message_ack is not None:
tool_messages.append(tool_message_ack)
if planned_step is not None:
tool_flow_steps.append(planned_step)
current_step = next_step
planned_actions.append(action_summary)
else:
# Format and enrich tool query for display before execution
tool_query = format_tool_query_for_display(tool_name, tool_args, combined_tool_query)
tool_query = await enrich_tool_query_for_display(tool_name, tool_args, tool_query)
# Yield execution reasoning block with enhanced query before execution
## change point: "πspecial reasoning blockπ:
stage = "retrieval_tool_execution"
user_friendly_tool_name = tool_name_shown_to_user(tool_name)
content = ""
reasoning_chunk_dict = reasoning_dict_maker(
stage=stage, tool_name=user_friendly_tool_name, tool_query=tool_query, content=content
)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
# For non-action tools (retrieval), execute via helper
result = await execute_retrieval_tool_and_build_step(
tool_name=tool_name,
tool_id=tool_id,
tool_args=tool_args,
combined_tool_query=combined_tool_query,
combined_tools=combined_tools,
query_flow_store=query_flow_store,
current_step=current_step,
tool_query=tool_query,
)
(
flow_step,
tool_message,
next_step_exec,
tool_query,
execution_success,
execution_error,
) = result # type: ignore[assignment]
log.info(f"ChatID: {chat_id} - Tool {tool_name} result: {tool_message.content}")
if tool_message is not None:
tool_messages.append(tool_message)
# Format tool message content for user-friendly display (remove UUIDs, URLs, etc.)
## change point: "πspecial reasoning blockπ:
stage = "retrieval_tool_execution_message"
cleaned_content = format_tool_message_for_display(tool_message.content)
user_friendly_tool_name = tool_name_shown_to_user(tool_name)
content = f"{user_friendly_tool_name}'s result: {cleaned_content}\n\n"
reasoning_chunk_dict = reasoning_dict_maker(
stage=stage, tool_name=user_friendly_tool_name, tool_query=tool_query, content=content
)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
else:
tool_messages.append(tool_message)
if flow_step is not None:
tool_flow_steps.append(flow_step)
current_step = next_step_exec
# If clarification was requested, stop further planning
if clarification_requested:
break
# Add tool results to conversation and continue
messages.extend(tool_messages)
# Get next response from LLM
response = await llm_with_method_tools.ainvoke(messages)
messages.append(response)
# Stream planner reasoning for this iteration (if present and not a control token)
# Only stream as reasoning if there are more tool calls coming (not the final response)
# This represents the LLM's internal thinking/reasoning process
has_more_tool_calls = bool(getattr(response, "tool_calls", None))
if has_more_tool_calls:
try:
_iter_content = str(getattr(response, "content", "") or "").strip()
if _iter_content:
## change point: "πspecial reasoning blockπ:
stage = "planner_tool_selection"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content=_iter_content)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
except Exception:
pass
# update _has_tool_calls
_has_tool_calls = has_more_tool_calls
# Log planner decisions for this iteration
try:
if _has_tool_calls:
tool_calls_list = getattr(response, "tool_calls", []) or []
tool_names = []
for _tc in tool_calls_list:
if isinstance(_tc, dict):
tool_names.append(_tc.get("name", ""))
else:
tool_names.append(getattr(_tc, "name", ""))
reasoning_text = str(getattr(response, "content", "") or "").strip()
reason_preview = reasoning_text or "(none)"
log.info(f"ChatID: {chat_id} - Planner selected tools (iteration {iteration_count}): {tool_names}")
log.info(f"ChatID: {chat_id} - Planner reasoning (iteration {iteration_count}): {reason_preview}")
else:
# Log when no tool calls are made in iteration
reasoning_text = str(getattr(response, "content", "") or "").strip()
log.info(f"ChatID: {chat_id} - Planner selected tools (iteration {iteration_count}): []")
log.info(f"ChatID: {chat_id} - Planner reasoning (iteration {iteration_count}, no tools): {reasoning_text}")
except Exception as e:
log.warning(f"ChatID: {chat_id} - Failed to log planner decisions for iteration {iteration_count}: {e}")
# Record planner decision for this iteration (selected tools + reasoning)
try:
selection_step, next_step_sel = build_planner_tool_selection_step(response, current_step=current_step)
if selection_step is not None:
try:
ed = selection_step.setdefault("execution_data", {})
ed["iteration"] = iteration_count
except Exception:
pass
tool_flow_steps.append(selection_step)
current_step = next_step_sel
except Exception:
pass
# Update loop condition for next iteration
loop_condition = (_has_tool_calls or len(planned_actions) == 0) and iteration_count < max_iterations
# Log why the tool selection loop exited
if clarification_requested:
# Persist any steps recorded so far (including clarification)
if tool_flow_steps:
with contextlib.suppress(Exception):
async with get_streaming_db_session() as _subdb:
await _upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=tool_flow_steps,
db=_subdb,
)
# End stream without free-form content
yield "__FINAL_RESPONSE__"
return
elif iteration_count >= max_iterations:
log.warning(f"Tool selection loop exited due to maximum iterations ({max_iterations}) reached")
if len(planned_actions) == 0:
# Signal failure and stop before emitting any free-form LLM content
msg = "An unexpected error occurred. Please try again later."
yield msg
# Persist any retrieval steps captured so far
# Append planning summary before persisting steps
try:
step, next_step_ps = build_planner_summary_step(
current_step=current_step,
iteration_count=iteration_count,
planned_actions_count=len(planned_actions),
loop_warning_detected=loop_warning_detected,
tool_calls_count=len(tool_call_history),
)
tool_flow_steps.append(step)
current_step = next_step_ps
except Exception:
pass
if tool_flow_steps:
with contextlib.suppress(Exception):
async with get_streaming_db_session() as _subdb:
await _upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=tool_flow_steps,
db=_subdb,
)
# Do not include free-form content; end stream
yield f"__FINAL_RESPONSE__{msg}"
return
elif len(planned_actions) == 0:
# No actions planned - this is a retrieval-only query or unsupported feature request
# The LLM should have provided an answer/explanation in its content
log.info(f"ChatID: {chat_id} - No actions planned; streaming LLM's response as final answer")
# Stream the LLM's content as the final answer
if hasattr(response, "content") and response.content:
content = str(response.content).strip()
if content:
# Stream the content in chunks for better UX
async for stream_chunk in stream_content_in_chunks(content):
yield stream_chunk
yield "\n"
else:
# Fallback if no content provided
fallback_msg = "I wasn't able to find a suitable response. Could you please rephrase your request?"
yield fallback_msg
yield "\n"
content = fallback_msg
else:
# Fallback if no content attribute
fallback_msg = "I wasn't able to find a suitable response. Could you please rephrase your request?"
yield fallback_msg
yield "\n"
content = fallback_msg
# Record the flow steps
try:
step, next_step_ps = build_planner_summary_step(
current_step=current_step,
iteration_count=iteration_count,
planned_actions_count=0,
loop_warning_detected=loop_warning_detected,
tool_calls_count=len(tool_call_history),
)
tool_flow_steps.append(step)
current_step = next_step_ps
except Exception:
pass
if tool_flow_steps:
with contextlib.suppress(Exception):
async with get_streaming_db_session() as _subdb:
await _upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=tool_flow_steps,
db=_subdb,
)
yield f"__FINAL_RESPONSE__{content}"
return
else:
log.debug(f"Tool selection loop completed successfully after {iteration_count} iterations with {len(planned_actions)} planned actions")
# Build final response and optional planner_no_actions step
final_response_chunks, current_step, no_actions_step = await build_and_stream_final_response(
response=response,
current_step=current_step,
iteration_count=iteration_count,
chat_id=chat_id,
)
if no_actions_step is not None:
tool_flow_steps.append(no_actions_step)
# Stream final response chunks - stream actual content
for chunk in final_response_chunks:
# Check if this is the actual response content
if hasattr(response, "content") and response.content:
content = str(response.content).strip()
# Stream actual content
if content:
chunk_without_newline = chunk.rstrip("\n")
# Check if this chunk matches the actual content (it should be content + "\n")
if chunk_without_newline == content:
# This is the actual response content - stream it in chunks
async for stream_chunk in stream_content_in_chunks(chunk_without_newline):
yield stream_chunk
# Add newline if it was in the original chunk
if chunk.endswith("\n"):
yield "\n"
continue
# Other message types - yield directly
yield chunk
# Record tool executions in database
# Append planning summary before persisting steps
try:
step, next_step_ps = build_planner_summary_step(
current_step=current_step,
iteration_count=iteration_count,
planned_actions_count=len(planned_actions),
loop_warning_detected=loop_warning_detected,
tool_calls_count=len(tool_call_history),
)
tool_flow_steps.append(step)
current_step = next_step_ps
except Exception:
pass
if tool_flow_steps:
async with get_streaming_db_session() as _subdb:
flow_step_result = await _upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=tool_flow_steps,
db=_subdb,
)
if flow_step_result["message"] != "success":
log.warning("Failed to record action execution in database")
# Return the complete response for storage in chat history
if final_response_chunks:
# Yield a special signal that can be detected by the calling function
# This follows the same pattern as response_processor.py
final_response = "".join(final_response_chunks)
yield f"__FINAL_RESPONSE__{final_response}"
else:
# Signal end without free-form content
yield "__FINAL_RESPONSE__"
except Exception:
# Log full traceback for diagnosis
log.error(f"ChatID: {chat_id} - Error in action execution", exc_info=True)
# Emit user-friendly error and final response so caller can persist
error_msg = "An unexpected error occurred. Please try again later."
try:
yield error_msg
yield f"__FINAL_RESPONSE__{error_msg}"
except Exception:
# Best effort; if streaming breaks here, nothing else we can do
pass
return
apps/pi/pi/services/chat/askmode_tool_executor.py
+459
View File
@@ -0,0 +1,459 @@
"""Retrieval tool execution logic for single and multi-tool scenarios.
Also includes tools for ask mode."""
import json
from collections.abc import AsyncIterator
from typing import Any
from typing import Dict
from typing import List
from typing import Tuple
from typing import Union
from langchain_core.messages import BaseMessage
from langchain_core.messages import HumanMessage
from langchain_core.messages import SystemMessage
from langchain_core.messages import ToolMessage
from pi import logger
from pi.app.models.enums import MessageMetaStepType
from pi.services.chat.helpers import ask_mode_helpers
from pi.services.chat.helpers.flow_tracking import check_and_build_clarification
from pi.services.chat.helpers.flow_tracking import execute_tool_step
from pi.services.chat.helpers.flow_tracking import orchestrate_llm_step
from pi.services.chat.helpers.flow_tracking import store_and_format_clarification
from pi.services.chat.helpers.tool_utils import format_tool_message_for_display
from pi.services.chat.helpers.tool_utils import format_tool_query_for_display
from pi.services.chat.helpers.tool_utils import stream_content_in_chunks
from pi.services.chat.helpers.tool_utils import tool_name_shown_to_user
from pi.services.chat.prompts import pai_ask_system_prompt
from pi.services.chat.utils import reasoning_dict_maker
log = logger.getChild(__name__)
def extract_facts_from_tool_result(tool_name: str, result: Any) -> Dict[str, Any]:
"""Extract structured facts from tool results for enhanced history."""
facts = {}
try:
# Handle StandardAgentResponse dict format (most common case)
# Tool handlers return dicts like: {"results": "...", "entity_urls": [...], ...}
if isinstance(result, dict):
# Extract entity URLs (critical for citations and linking)
entity_urls = result.get("entity_urls")
if entity_urls and isinstance(entity_urls, list):
facts["entity_urls"] = entity_urls[:10] # Limit to prevent bloat
facts["entity_urls_count"] = len(entity_urls)
# Extract IDs from entity URLs for easier reference
entity_ids = [url.get("id") for url in entity_urls if isinstance(url, dict) and url.get("id")]
if entity_ids:
facts["entity_ids"] = entity_ids[:10]
# Extract intermediate results (SQL debugging for structured_db_tool)
intermediate_results = result.get("intermediate_results")
if intermediate_results and isinstance(intermediate_results, dict):
if intermediate_results.get("generated_sql"):
facts["sql_query"] = str(intermediate_results["generated_sql"])[:300]
if intermediate_results.get("relevant_tables"):
facts["relevant_tables"] = intermediate_results["relevant_tables"]
if intermediate_results.get("final_query"):
facts["final_query"] = str(intermediate_results["final_query"])[:200]
# Extract search metadata (for vector/semantic search tools)
execution_metadata = result.get("execution_metadata")
if execution_metadata and isinstance(execution_metadata, dict):
facts["search_metadata"] = execution_metadata
# Extract result preview from "results" field
results_text = result.get("results")
if results_text:
facts["result_preview"] = str(results_text)[:300]
# Extract other common fields that might be in the dict
for key in ["entity_id", "entity_name", "entity_url", "count", "ids", "results_count"]:
if key in result and result[key]:
facts[key] = result[key]
# Handle string results (simple tool responses)
elif isinstance(result, str):
facts["result_preview"] = result[:300]
except Exception as e:
log.debug(f"Could not extract facts from {tool_name}: {e}")
return facts
async def execute_tools_for_ask_mode(
chatbot_instance,
user_meta,
workspace_id,
workspace_slug,
project_id,
conversation_history,
enhanced_conversation_history, # 🆕 Enhanced context with action details
user_id,
chat_id,
query_flow_store,
enhanced_query_for_processing, # This is what's actually passed from chat.py
query_id,
step_order,
db,
parsed_query,
reasoning_container=None,
) -> AsyncIterator[Union[str, Dict[str, Any]]]:
"""Execute tools for ask mode"""
log.info(f"ChatID: {chat_id} - Executing tools for ask mode with enhanced_conversation_history: {enhanced_conversation_history}")
try:
# Extract workspace_in_context from query_flow_store
workspace_in_context = True
if isinstance(query_flow_store, dict):
workspace_in_context = query_flow_store.get("workspace_in_context", True)
log.info(f"ChatID: {chat_id} - Workspace in context: {workspace_in_context}")
# Resolve workspace_id and workspace_slug (only when workspace is in context)
workspace_id, workspace_slug = await ask_mode_helpers.resolve_workspace_id_and_slug(
workspace_id=workspace_id,
project_id=project_id,
chat_id=chat_id,
)
tools = await chatbot_instance._create_tools_for_ask_mode(
db,
user_meta,
workspace_id,
workspace_slug,
project_id,
user_id,
chat_id,
query_flow_store,
conversation_history,
query_id,
workspace_in_context=workspace_in_context,
chatbot_instance=chatbot_instance,
)
if not tools:
log.warning("Unable to initialize tools. Please try again.")
msg = "An unexpected error occurred. Please try again later."
yield msg
return
try:
# Determine if this is a clarification follow-up (raw user input)
is_clarification_followup = bool(user_meta and isinstance(user_meta, dict) and user_meta.get("clarification_context"))
query_label = "Query (raw user input from clarification)" if is_clarification_followup else "Query"
# Log comprehensive debugging information
log.info(f"ChatID: {chat_id} - {query_label}: {enhanced_query_for_processing}")
log.info(f"ChatID: {chat_id} - Ask mode LLM input - Available tools: {[t.name for t in tools]}")
except Exception as e:
log.warning(f"ChatID: {chat_id} - Failed to log debug info: {e}")
# Build enhanced prompt and context block
custom_prompt, context_block = await ask_mode_helpers.construct_enhanced_prompt_and_context(
enhanced_conversation_history=enhanced_conversation_history,
project_id=project_id,
workspace_id=workspace_id,
user_id=user_id,
user_meta=user_meta,
workspace_in_context=workspace_in_context,
)
# Bind tools to the LLM - explicitly set tool_choice to 'auto' to ensure tool use is enabled
# deduplicate tools
tools = list({tool.name: tool for tool in tools}.values())
llm_with_tools = chatbot_instance.tool_llm.bind_tools(tools)
# Set tracking context on the bound LLM instance
llm_with_tools.set_tracking_context(query_id, db, MessageMetaStepType.TOOL_ORCHESTRATION, chat_id=str(chat_id))
# Track execution - use FlowStepCollector context manager for automatic persistence
current_step = step_order
messages: List[BaseMessage] = [] # type: ignore[no-redef]
responses: List[Tuple[str, str, Union[str, Dict[str, Any]]]] = []
system_prompt_to_use = f"{pai_ask_system_prompt}\n\n{context_block}"
messages.append(SystemMessage(content=system_prompt_to_use))
query_to_use = f"{custom_prompt}\n\nUser Query: {enhanced_query_for_processing}"
messages.append(HumanMessage(content=query_to_use))
# Log what's sent to the tool selection LLM
log.info(f"ChatID: {chat_id} - Ask mode LLM input - Query: {query_to_use}")
## change point: "πspecial reasoning blockπ:
# Yield initial status - different message for clarification follow-ups
if is_clarification_followup:
stage = "ask_mode_clarification_followup"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
else:
stage = "ask_mode_beginning"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
# Initial invocation of LLM with tools (decorator persists selection/context/reasoning)
ai_message, current_step = await orchestrate_llm_step(
llm_with_tools,
messages,
query_id=query_id,
chat_id=chat_id,
db=db,
current_step=current_step,
include_context=True,
enhanced_conversation_history=enhanced_conversation_history,
enhanced_query_for_processing=enhanced_query_for_processing,
tools=tools,
)
# Log LLM's initial response for debugging
log.info(f"ChatID: {chat_id} - Has Tool Calls: {hasattr(ai_message, "tool_calls") and bool(getattr(ai_message, "tool_calls", None))}")
# Check if LLM made any tool calls
if not ai_message.tool_calls:
log.info(f"ChatID: {chat_id} - No tool calls found in the initial LLM response.")
stage = "final_response"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
if ai_message.content:
# Stream the final answer in chunks to simulate streaming
async for chunk in stream_content_in_chunks(str(ai_message.content)):
yield chunk
return
# Stream LLM's initial reasoning content (internal thinking)
# Only stream as reasoning if there ARE tool calls (otherwise it's the final answer)
try:
if hasattr(ai_message, "content") and ai_message.content:
reasoning_content = str(ai_message.content).strip()
if reasoning_content:
stage = "planner_tool_selection"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content=reasoning_content)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
except Exception:
pass
# Add the AIMessage with tool_calls to conversation BEFORE processing tool calls
messages.append(ai_message)
# Process tool calls iteratively; persist steps via decorators
while ai_message.tool_calls:
# Execute all tool calls
tool_messages: List[ToolMessage] = []
for tool_call in ai_message.tool_calls:
tool_name = tool_call["name"]
tool_args = tool_call["args"]
tool_id = tool_call["id"]
# Intercept clarification requests and short-circuit
if tool_name == "ask_for_clarification":
log.info(f"ChatID: {chat_id} - Clarification requested by LLM: {tool_args.get("reason", "No reason provided")}")
clarification_payload, result = await check_and_build_clarification(tools=tools, tool_name=tool_name, tool_args=tool_args)
formatted_text, current_step = await store_and_format_clarification(
query_id=query_id,
chat_id=chat_id,
db=db,
current_step=current_step,
tool_name=tool_name,
tool_args=tool_args,
result=result,
clarification_payload=clarification_payload,
enhanced_query_for_processing=enhanced_query_for_processing,
tools=tools,
)
yield formatted_text
return
else:
# Find and execute the tool (decorator persists success/error)
tool_to_execute = next((t for t in tools if t.name == tool_name), None)
if tool_to_execute:
log.info(f"ChatID: {chat_id} - Executing tool: {tool_name} with args: {str(tool_args)[:100]}")
try:
# Yield tool execution status
user_friendly_tool_name = tool_name_shown_to_user(tool_name)
log.info(f"ChatID: {chat_id} - Tool name: {tool_name} - User friendly tool name: {user_friendly_tool_name}")
# Format tool query for display (same as build mode)
tool_query_str = format_tool_query_for_display(tool_name, tool_args, enhanced_query_for_processing)
stage = "retrieval_tool_execution"
content = ""
reasoning_chunk_dict = reasoning_dict_maker(
stage=stage, tool_name=user_friendly_tool_name, tool_query=tool_query_str, content=content
)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
# Provide current step_order to downstream structured DB tool sub-steps
try:
if tool_name == "structured_db_tool" and isinstance(query_flow_store, dict):
query_flow_store["step_order"] = current_step
except Exception:
pass
# Execute the tool with persistence
tool_result, current_step = await execute_tool_step(
tool_to_execute,
tool_name,
tool_args,
query_id=query_id,
chat_id=chat_id,
db=db,
current_step=current_step,
extract_facts_fn=extract_facts_from_tool_result,
)
# Log brief output for observability
result_preview = str(tool_result) if tool_result else "None"
log.info(f"ChatID: {chat_id} - Tool {tool_name} result preview: {result_preview[:200]}")
log.debug(f"ChatID: {chat_id} - Tool {tool_name} completed successfully")
# Track the response
responses.append((tool_name, str(tool_args), tool_result))
# Create tool message for LLM
tool_message = ToolMessage(content=str(tool_result), tool_call_id=str(tool_id))
tool_messages.append(tool_message)
# Format tool result for display in reasoning block
if isinstance(tool_result, dict):
# Format as "message\n\nResult: {data}" so format_tool_message_for_display can clean it for users
# while the LLM gets the full structured data with UUIDs
message = tool_result.get("message", "")
# If there's a 'data' field, use it; otherwise omit the Result section (simpler format)
if "data" in tool_result and tool_result["data"]:
tool_content = f"{message}\n\nResult: {json.dumps(tool_result["data"], ensure_ascii=False)}"
else:
# No data field, just use the message
tool_content = message
else:
tool_content = str(tool_result)
# Format and stream tool message content for user-friendly display (remove UUIDs, URLs, etc.)
cleaned_content = format_tool_message_for_display(tool_content)
if cleaned_content and cleaned_content.strip():
stage = "retrieval_tool_execution_message"
content = f"{user_friendly_tool_name}'s result: {cleaned_content}\n\n"
reasoning_chunk_dict = reasoning_dict_maker(
stage=stage, tool_name=user_friendly_tool_name, tool_query=tool_query_str, content=content
)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
except Exception as e:
log.error(f"ChatID: {chat_id} - Error executing tool {tool_name}: {str(e)}")
# Yield error status
stage = "tool_error"
error_status_message = f"⚠️ Error executing {user_friendly_tool_name}\n\n"
reasoning_chunk_dict = reasoning_dict_maker(
stage=stage,
tool_name=user_friendly_tool_name,
tool_query=tool_query_str if "tool_query_str" in locals() else "",
content=error_status_message,
)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
# Error step already persisted by decorator; step updated
error_tool_message = ToolMessage(content=f"Error executing {tool_name}: {str(e)}", tool_call_id=str(tool_id or ""))
tool_messages.append(error_tool_message)
else:
log.error(f"ChatID: {chat_id} - Tool {tool_name} not found")
stage = "tool_unavailable"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
error_tool_message = ToolMessage(content=f"Tool {tool_name} not found", tool_call_id=str(tool_id or ""))
tool_messages.append(error_tool_message)
# Add tool messages to conversation
messages.extend(tool_messages)
# Yield status before getting next LLM response
stage = "ask_mode_analyzing_results"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
# Get next response from LLM (persist follow-up selection/reasoning)
ai_message, current_step = await orchestrate_llm_step(
llm_with_tools,
messages,
query_id=query_id,
chat_id=chat_id,
db=db,
current_step=current_step,
include_context=False,
enhanced_conversation_history=None,
enhanced_query_for_processing=None,
tools=tools,
)
# Handle failure case
if ai_message == "TOOL_ORCHESTRATION_FAILURE":
stage = "final_response"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
yield "An unexpected error occurred. Please try again later."
return
# Stream LLM's reasoning content after tool execution (internal thinking)
# Only stream as reasoning if there are more tool calls coming (not the final response)
has_more_tool_calls = bool(getattr(ai_message, "tool_calls", None))
if has_more_tool_calls:
try:
if hasattr(ai_message, "content") and ai_message.content:
reasoning_content = str(ai_message.content).strip()
if reasoning_content:
stage = "planner_tool_selection"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content=reasoning_content)
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
except Exception:
pass
messages.append(ai_message)
# Follow-up selection persisted by decorator above
# Final response from LLM (no more tool calls)
log.info(f"ChatID: {chat_id} - Tool orchestration complete, generating final response")
# Yield final response generation status
stage = "final_response"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
# Stream the final answer back to caller
if ai_message.content:
final_content = str(ai_message.content)
log.info(f"ChatID: {chat_id} - Final response length: {len(final_content)} chars")
# Stream the final answer in chunks to simulate streaming
async for chunk in stream_content_in_chunks(final_content):
yield chunk
else:
log.warning(f"ChatID: {chat_id} - LLM provided no final content")
yield "I've gathered the information but couldn't generate a response. Please try again."
except Exception as e:
log.error(f"ChatID: {chat_id} - Error in execute_tools_for_ask_mode: {str(e)}")
raise e
apps/pi/pi/services/chat/chat.py
+265 -783
View File
File diff suppressed because it is too large Load Diff
apps/pi/pi/services/chat/helpers/action_execution_helpers.py
+421
View File
@@ -0,0 +1,421 @@
"""Helper functions for batch action execution."""
import json
from datetime import datetime
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from uuid import UUID
from sqlalchemy import select
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.models.message import Message
from pi.app.models.message import MessageFlowStep
from pi.app.schemas.chat import ArtifactData
from pi.services.chat.helpers.build_mode_helpers import TOOL_NAME_TO_CATEGORY_MAP
from pi.services.chat.helpers.tool_utils import classify_tool
from pi.services.chat.prompts import plane_context
from pi.services.retrievers.pg_store.action_artifact import get_action_artifacts_by_ids
log = logger.getChild(__name__)
async def load_artifacts(request_data: List[ArtifactData], db: AsyncSession) -> Tuple[List[Dict], str, str]:
"""Load and validate artifacts for execution."""
artifacts = await get_action_artifacts_by_ids(db, [a.artifact_id for a in request_data])
original_query = artifacts[0].data.get("planning_context", {}).get("original_query", "")
conversation_context = artifacts[0].data.get("planning_context", {}).get("conversation_context", {})
planned_actions = []
for artifact, req_item in zip(artifacts, request_data):
if req_item.is_edited:
tool_args = req_item.action_data
# TODO: validate tool args
else:
tool_args = artifact.data.get("tool_args_raw", {})
planned_actions.append({
"artifact_id": str(artifact.id),
"tool_name": artifact.data.get("planning_data", {}).get("tool_name", ""),
"args": tool_args,
"entity_type": artifact.data.get("planning_data", {}).get("artifact_type", ""),
"action": artifact.data.get("planning_data", {}).get("action", ""),
})
return planned_actions, original_query, conversation_context
async def update_flow_steps(results, message_id, chat_id, db: AsyncSession):
"""Mark executed actions in database."""
from pi.services.retrievers.pg_store.action_artifact import update_action_artifact_execution_status
for r in results:
artifact_id = r.get("artifact_id")
if artifact_id:
artifact_uuid = UUID(artifact_id)
# Extract entity_id from entity_info if available
entity_info = r.get("entity_info")
entity_id = None
if entity_info and isinstance(entity_info, dict):
# Try to get entity_id from entity_info
entity_id_str = entity_info.get("entity_id")
if entity_id_str:
try:
entity_id = UUID(entity_id_str)
except (ValueError, TypeError):
log.warning(f"Invalid entity_id format in entity_info: {entity_id_str}")
# Update ActionArtifact and MessageFlowStep execution status
await update_action_artifact_execution_status(
db=db,
message_id=message_id,
chat_id=chat_id,
artifact_id=artifact_uuid,
is_executed=True,
success=r.get("success", False),
entity_id=entity_id,
entity_info=entity_info,
execution_result=r.get("result", ""),
executed_at=r.get("executed_at"),
)
def format_response(planned_actions, results, start_time) -> Dict[str, Any]:
"""Format the execution response with clean, non-redundant structure."""
log.info(f"\n\nResults in format_response: {results}\n")
try:
total_planned = len(planned_actions)
completed_count = sum(1 for r in results if r.get("success"))
failed_count = sum(1 for r in results if r.get("success") is False)
response: Dict[str, Any] = {
"action_summary": {
"total_planned": total_planned,
"completed": completed_count,
"failed": failed_count,
"duration_seconds": round((datetime.utcnow() - start_time).total_seconds(), 2),
},
}
if results:
response["actions"] = create_clean_actions_response(results)
log.info(f"\n\nRESPONSE: {response}\n\n")
return response
except Exception as e:
log.error(f"Error formatting execution response: {e}")
return {
"action_summary": {
"total_planned": len(planned_actions),
"completed": 0,
"failed": len(planned_actions),
"duration_seconds": round((datetime.utcnow() - start_time).total_seconds(), 2),
},
"actions": [],
}
async def get_planned_actions_for_execution(message_id: UUID, chat_id: UUID, db: AsyncSession) -> List[Dict[str, Any]]:
"""Retrieve all planned actions for a message that are ready for execution."""
try:
# Get flow steps with planned actions - filter for TOOL steps that are planned
stmt = (
select(MessageFlowStep)
.where(MessageFlowStep.message_id == message_id) # type: ignore[arg-type]
.where(MessageFlowStep.chat_id == chat_id) # type: ignore[arg-type]
.where(MessageFlowStep.is_planned == True) # type: ignore[arg-type] # noqa: E712
.where(MessageFlowStep.is_executed == False) # type: ignore[arg-type] # noqa: E712
.order_by(MessageFlowStep.step_order) # type: ignore[arg-type]
)
result = await db.execute(stmt)
flow_steps = result.scalars().all()
planned_actions = []
for step in flow_steps:
execution_data = step.execution_data or {}
planned_action = {
"step_id": str(step.id),
"step_order": step.step_order,
"tool_name": step.tool_name,
"args": execution_data.get("args", {}),
"action_summary": execution_data.get("action_summary", {}),
"tool_id": execution_data.get("tool_id"),
"artifact_id": execution_data.get("artifact_id"), # NEW: Include artifact ID
"sequence": step.step_order, # NEW: Include planned sequence to propagate
}
planned_actions.append(planned_action)
return planned_actions
except Exception as e:
log.error(f"Error retrieving planned actions: {e}")
return []
async def get_original_user_query(message_id: UUID, db: AsyncSession) -> Optional[str]:
"""Get the original user query for the message."""
try:
stmt = select(Message).where(Message.id == message_id) # type: ignore[arg-type]
result = await db.execute(stmt)
message = result.scalar_one_or_none()
if message:
return message.content
else:
log.warning(f"Message {message_id} not found")
return None
except Exception as e:
log.error(f"Error retrieving original user query: {e}")
return None
def build_execution_prompt(original_query: str, planned_actions: List[Dict[str, Any]], conversation_context: Dict[str, Any]) -> str:
"""Build the execution prompt for the LLM."""
# Format planned actions for the prompt
actions_text = ""
tool_categories = set()
for i, action in enumerate(planned_actions, 1):
action_summary = action.get("action_summary", {})
tool_name = action.get("tool_name", "")
args = action.get("args", {})
entity_type = action.get("entity_type", "")
if entity_type:
tool_categories.add(entity_type)
# Create a readable description of the action
action_desc = action_summary.get("action", tool_name)
action_summary.get("details", "")
actions_text += f"{i}. {action_desc}\n"
actions_text += f" Tool: {tool_name}\n"
actions_text += f" Arguments: {json.dumps(args, indent=4)}\n\n"
# Format conversation context for the prompt
conversation_context_text = ""
for key, value in conversation_context.items():
conversation_context_text += f"{key}: {value}\n"
execution_prompt = f"""You are an API tool calling agent for Plane. You are executing approved actions for this user request: "{original_query}"
context about Plane:
{plane_context}
PLANNED ACTIONS (approved by user):
{actions_text}
EXECUTION MODE: You must now EXECUTE these actions in the correct order using the available tools.
EXECUTION GUIDELINES:
- Execute each planned action exactly once using the provided tool and arguments
- Use outputs from earlier actions as inputs to later actions when needed
- If an action creates an entity (like an issue, cycle, etc.), use the returned ID in subsequent actions that need it
- Execute actions in logical dependency order (e.g., create issue first, then add to cycle)
- If an action fails, stop execution and report the specific error
- Provide clear status updates after each successful action
**CRITICAL: PLACEHOLDER HANDLING**
- Some arguments may contain placeholders that need to be resolved
- You should use the UUIDs, names, identifiers, etc. from the results returned by the previous actions to resolve the placeholders for the next action.
- Careful with the placeholder types. Use the tool description to determine the type of the placeholder, and then determine whether to use the UUID, name, identifier, etc.
Entity Search Tools:
In addition to the modifying tools provided to you, you also have access to entity search tools (search_project_by_identifier, search_user_by_name, etc.). Use them ONLY if:
1. An argument is completely MISSING (e.g., no project_id at all)
2. An argument contains an EXPLICIT placeholder like: "<id of project: ProjectName>"
3. An argument has a non-UUID value that needs resolution (e.g., project_id="PROJ")
CRITICAL: You must use the EXACT tool names provided above. Do not modify, shorten, or change the tool names in any way.
EXAMPLE: If the planned action shows tool_name: "workitems_create", you must call "workitems_create" exactly, not "create_workitem" or "workitem_create".
IMPORTANT: Start executing the actions now. Use the exact tool names and arguments provided in the planned actions above.""" # noqa: E501
return execution_prompt
def has_placeholders(tool_args: Dict[str, Any]) -> bool:
"""Check if tool arguments contain placeholders or non-UUID *_id values that need resolution."""
import re
# Regex for matching standard UUIDs - fixed pattern to handle UUIDs properly
uuid_regex = re.compile(r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")
# Include both "workitem" and "issue" as aliases for the same entity
entity_keys = {"module", "workitem", "issue", "project", "cycle", "label", "state", "user", "assignee"}
for key, value in tool_args.items():
# Special-case: workspace scope sentinel should NOT trigger resolution
if key == "project_id" and isinstance(value, str) and value == "__workspace_scope__":
continue
# 1) Explicit placeholder syntax
if isinstance(value, str) and "<id of" in value:
return True
# 2) Non-UUID string passed to an *_id field for known entities
if isinstance(value, str) and (key.endswith("_id") or key.endswith("_ids")):
# Handle both _id and _ids
if key.endswith("_ids"):
entity_type = key[:-4] # remove _ids
else:
entity_type = key[:-3] # remove _id
# Treat issue_id as workitem for resolution purposes
if entity_type == "issue":
entity_type = "workitem"
# For single value, check if it matches UUID
if entity_type in entity_keys:
is_uuid = uuid_regex.match(value)
if not is_uuid:
return True
# 3) List of placeholders or non-UUID id values for known entities
if isinstance(value, list):
for item in value:
# Skip sentinel in lists as well (defensive)
if key == "project_id" and isinstance(item, str) and item == "__workspace_scope__":
continue
if isinstance(item, str):
if "<id of" in item:
return True
# Check for non-UUIDs in _id or _ids fields
if key.endswith("_id") or key.endswith("_ids"):
if key.endswith("_ids"):
entity_type = key[:-4]
else:
entity_type = key[:-3]
if entity_type == "issue":
entity_type = "workitem"
if entity_type in entity_keys and not uuid_regex.match(item):
return True
return False
def create_clean_actions_response(executed_actions: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Create clean action results without duplicate entity information.
Filters out retrieval/search tools and only includes actual modifying actions.
"""
clean_actions = []
log.info(f"\n\nEXECUTED ACTIONS in create_clean_actions_response: {executed_actions}\n\n")
for action in executed_actions:
# Extract action from tool name (e.g., "workitems_create" -> "create")
tool_name = action.get("tool_name", "")
# Filter out retrieval tools - only show actual actions to frontend
if tool_name:
is_retrieval, is_action = classify_tool(tool_name)
if is_retrieval and not is_action:
log.debug(f"Filtering out retrieval tool from actions response: {tool_name}")
continue
action_type = TOOL_NAME_TO_CATEGORY_MAP.get(tool_name, {}).get("action_type", "")
action_data = {
"action": action_type,
"artifact_type": action.get("artifact_type"),
"success": action.get("success"),
"executed_at": action.get("executed_at"),
"artifact_id": action.get("artifact_id"), # NEW: Include artifact ID
"sequence": action.get("sequence"), # NEW: Include planned step order
"version_number": action.get("version_number"), # NEW: Include version sequence number
}
if action.get("success"):
# For successful actions, include essential entity info
entity_info = action.get("entity_info")
if entity_info and isinstance(entity_info, dict):
# Only include the most important entity fields
essential_entity = {}
for field in ["entity_url", "entity_name", "entity_type", "entity_id"]:
if field in entity_info and entity_info[field]:
essential_entity[field] = entity_info[field]
# Include issue_identifier when available (for work-items)
if entity_info.get("issue_identifier"):
essential_entity["issue_identifier"] = entity_info["issue_identifier"]
if essential_entity:
action_data["entity"] = essential_entity
# Add project_identifier at the action root when derivable
project_identifier = None
try:
# Prefer deriving from identifiers present in entity_info
if entity_info.get("entity_type") == "project":
project_identifier = entity_info.get("entity_identifier") or entity_info.get("project_identifier")
elif entity_info.get("issue_identifier") and isinstance(entity_info.get("issue_identifier"), str):
ident = entity_info.get("issue_identifier")
if "-" in ident:
project_identifier = ident.split("-", 1)[0]
elif entity_info.get("entity_url") and isinstance(entity_info.get("entity_url"), str):
url = entity_info.get("entity_url")
# Attempt to parse /browse/PROJECT-SEQ/ pattern
if "/browse/" in url:
after = url.split("/browse/", 1)[1]
ident = after.split("/", 1)[0]
if "-" in ident:
project_identifier = ident.split("-", 1)[0]
except Exception:
project_identifier = None
if project_identifier:
action_data["project_identifier"] = project_identifier
# Extract the nice success message from the result
result = action.get("result", "")
if result:
# Extract the nice message that comes after "✅ " and before "\n\n"
nice_message = extract_success_message(result)
if nice_message:
action_data["message"] = nice_message
else:
# Fallback to generic messages
if "created" in result.lower():
action_data["message"] = "Created successfully"
elif "updated" in result.lower():
action_data["message"] = "Updated successfully"
else:
action_data["message"] = "Action completed successfully"
else:
# For failed actions, include error message
error = action.get("error", "")
if error:
action_data["error"] = error[:200] + "..." if len(error) > 200 else error
clean_actions.append(action_data)
return clean_actions
def extract_success_message(result: str) -> str:
"""Extract the nice success message from the tool result."""
if not result or not isinstance(result, str):
return ""
# Look for the pattern: "✅ [message]\n\n"
if "✅" in result:
lines = result.split("\n")
for line in lines:
line = line.strip()
if line.startswith("✅"):
# Remove the "✅ " prefix and return the message
message = line[2:].strip() # Remove "✅ " (2 characters)
return message
return ""
apps/pi/pi/services/chat/helpers/action_executor.py
-1876
View File
File diff suppressed because it is too large Load Diff
apps/pi/pi/services/chat/helpers/action_property_mapper.py
+588
View File
@@ -0,0 +1,588 @@
"""
Helper module to map flat tool arguments to the nested 'properties' structure
required for the action summary in the UI.
"""
import logging
import re
from typing import Any
from typing import Dict
log = logging.getLogger(__name__)
def is_uuid(value: str) -> bool:
"""Check if a string is a valid UUID."""
uuid_pattern = re.compile(r"^[0-9a-f]{8}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{12}$", re.IGNORECASE)
return bool(uuid_pattern.match(value.strip()))
def map_workitem_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for workitem creation/update tools."""
properties: Dict[str, Any] = {}
# Priority: {"name": "low"}
if "priority" in tool_args and tool_args["priority"]:
properties["priority"] = {"name": tool_args["priority"]}
# State: {"id": "uuid"} or {"name": "Done"}
if "state" in tool_args and tool_args["state"]:
state_val = tool_args["state"]
if is_uuid(state_val):
properties["state"] = {"id": state_val}
else:
properties["state"] = {"name": state_val}
# Assignees: [{"id": "uuid"}, ...]
if "assignees" in tool_args and tool_args["assignees"]:
assignees = tool_args["assignees"]
if isinstance(assignees, list):
properties["assignees"] = [{"id": assignee_id} for assignee_id in assignees if assignee_id]
# Labels: [{"id": "uuid"}, ...]
if "labels" in tool_args and tool_args["labels"]:
labels = tool_args["labels"]
if isinstance(labels, list):
properties["labels"] = [{"id": label_id} for label_id in labels if label_id]
# Dates: {"name": "YYYY-MM-DD"}
if "start_date" in tool_args and tool_args["start_date"]:
properties["start_date"] = {"name": tool_args["start_date"]}
if "target_date" in tool_args and tool_args["target_date"]:
properties["target_date"] = {"name": tool_args["target_date"]}
# Type ID: {"id": "uuid"}
if "type_id" in tool_args and tool_args["type_id"]:
properties["type_id"] = {"id": tool_args["type_id"]}
# Parent: {"id": "uuid"}
if "parent" in tool_args and tool_args["parent"]:
properties["parent"] = {"id": tool_args["parent"]}
# External ID: {"name": "external-123"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
# External Source: {"name": "jira"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
# Is Draft: {"name": true/false}
if "is_draft" in tool_args and tool_args["is_draft"] is not None:
properties["is_draft"] = bool(tool_args["is_draft"])
return properties
def map_project_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for project creation/update tools."""
properties: Dict[str, Any] = {}
# Identifier: {"name": "PROJ"}
if "identifier" in tool_args and tool_args["identifier"]:
properties["identifier"] = {"name": tool_args["identifier"]}
# Description: {"name": "text"} - UI often expects simple text wrapped
if "description" in tool_args and tool_args["description"]:
properties["description"] = {"name": tool_args["description"]}
# Project Lead: {"id": "uuid"}
if "project_lead" in tool_args and tool_args["project_lead"]:
properties["project_lead"] = {"id": tool_args["project_lead"]}
# Default Assignee: {"id": "uuid"}
if "default_assignee" in tool_args and tool_args["default_assignee"]:
properties["default_assignee"] = {"id": tool_args["default_assignee"]}
# Icon Prop: {"name": {...}} - JSON object
if "icon_prop" in tool_args and tool_args["icon_prop"]:
properties["icon_prop"] = {"name": tool_args["icon_prop"]}
# Emoji: {"name": "🚀"}
if "emoji" in tool_args and tool_args["emoji"]:
properties["emoji"] = {"name": tool_args["emoji"]}
# Cover Image: {"name": "url"}
if "cover_image" in tool_args and tool_args["cover_image"]:
properties["cover_image"] = {"name": tool_args["cover_image"]}
# Module View: {"name": true/false}
if "module_view" in tool_args and tool_args["module_view"] is not None:
properties["module_view"] = bool(tool_args["module_view"])
# Cycle View: {"name": true/false}
if "cycle_view" in tool_args and tool_args["cycle_view"] is not None:
properties["cycle_view"] = bool(tool_args["cycle_view"])
# Issue Views View: {"name": true/false}
if "issue_views_view" in tool_args and tool_args["issue_views_view"] is not None:
properties["issue_views_view"] = bool(tool_args["issue_views_view"])
# Page View: {"name": true/false}
if "page_view" in tool_args and tool_args["page_view"] is not None:
properties["page_view"] = bool(tool_args["page_view"])
# Intake View: {"name": true/false}
if "intake_view" in tool_args and tool_args["intake_view"] is not None:
properties["intake_view"] = bool(tool_args["intake_view"])
# Guest View All Features: {"name": true/false}
if "guest_view_all_features" in tool_args and tool_args["guest_view_all_features"] is not None:
properties["guest_view_all_features"] = bool(tool_args["guest_view_all_features"])
# Archive In: {"name": 12} - number of months
if "archive_in" in tool_args and tool_args["archive_in"] is not None:
properties["archive_in"] = int(tool_args["archive_in"])
# Close In: {"name": 12} - number of months
if "close_in" in tool_args and tool_args["close_in"] is not None:
properties["close_in"] = int(tool_args["close_in"])
# Timezone: {"name": "UTC"}
if "timezone" in tool_args and tool_args["timezone"]:
properties["timezone"] = {"name": tool_args["timezone"]}
# Time Tracking Enabled: {"name": true/false}
# Handle both variants: time_tracking_enabled and is_time_tracking_enabled
if "time_tracking_enabled" in tool_args and tool_args["time_tracking_enabled"] is not None:
properties["time_tracking_enabled"] = bool(tool_args["time_tracking_enabled"])
elif "is_time_tracking_enabled" in tool_args and tool_args["is_time_tracking_enabled"] is not None:
properties["is_time_tracking_enabled"] = bool(tool_args["is_time_tracking_enabled"])
# Issue Type Enabled: {"name": true/false}
if "is_issue_type_enabled" in tool_args and tool_args["is_issue_type_enabled"] is not None:
properties["is_issue_type_enabled"] = bool(tool_args["is_issue_type_enabled"])
return properties
def map_module_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for module creation/update tools and work item operations."""
properties: Dict[str, Any] = {}
# Start Date: {"name": "YYYY-MM-DD"}
if "start_date" in tool_args and tool_args["start_date"]:
properties["start_date"] = {"name": tool_args["start_date"]}
# Target Date: {"name": "YYYY-MM-DD"}
if "target_date" in tool_args and tool_args["target_date"]:
properties["target_date"] = {"name": tool_args["target_date"]}
# Status: {"name": "backlog/unstarted/started/completed/cancelled/paused"}
if "status" in tool_args and tool_args["status"]:
properties["status"] = {"name": tool_args["status"]}
# Lead: {"id": "uuid"}
if "lead" in tool_args and tool_args["lead"]:
properties["lead"] = {"id": tool_args["lead"]}
# Members: [{"id": "uuid"}, ...]
if "members" in tool_args and tool_args["members"]:
members = tool_args["members"]
if isinstance(members, list):
properties["members"] = [{"id": member_id} for member_id in members if member_id]
# Issues: [{"id": "uuid"}, ...] - for modules_add_work_items
if "issues" in tool_args and tool_args["issues"]:
issues = tool_args["issues"]
if isinstance(issues, list):
properties["issues"] = [{"id": issue_id} for issue_id in issues if issue_id]
# Issue ID: {"id": "uuid"} - for modules_remove_work_item
if "issue_id" in tool_args and tool_args["issue_id"]:
properties["issue_id"] = {"id": tool_args["issue_id"]}
# External ID: {"name": "external-123"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
# External Source: {"name": "jira"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
return properties
def map_cycle_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for cycle creation/update tools and work item operations."""
properties: Dict[str, Any] = {}
# Start Date: {"name": "YYYY-MM-DD"}
if "start_date" in tool_args and tool_args["start_date"]:
properties["start_date"] = {"name": tool_args["start_date"]}
# End Date: {"name": "YYYY-MM-DD"}
if "end_date" in tool_args and tool_args["end_date"]:
properties["end_date"] = {"name": tool_args["end_date"]}
# Owned By: {"id": "uuid"}
if "owned_by" in tool_args and tool_args["owned_by"]:
properties["owned_by"] = {"id": tool_args["owned_by"]}
# Issues: [{"id": "uuid"}, ...] - for cycles_add_work_items
if "issues" in tool_args and tool_args["issues"]:
issues = tool_args["issues"]
if isinstance(issues, list):
properties["issues"] = [{"id": issue_id} for issue_id in issues if issue_id]
# Issue ID: {"id": "uuid"} - for cycles_remove_work_item
if "issue_id" in tool_args and tool_args["issue_id"]:
properties["issue_id"] = {"id": tool_args["issue_id"]}
# New Cycle ID: {"id": "uuid"} - for cycles_transfer_work_items
if "new_cycle_id" in tool_args and tool_args["new_cycle_id"]:
properties["new_cycle_id"] = {"id": tool_args["new_cycle_id"]}
# External ID: {"name": "external-123"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
# External Source: {"name": "jira"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
# Timezone: {"name": "UTC"}
if "timezone" in tool_args and tool_args["timezone"]:
properties["timezone"] = {"name": tool_args["timezone"]}
return properties
def map_intake_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for intake creation/update tools."""
properties: Dict[str, Any] = {}
# Priority: {"name": "low/medium/high/urgent/none"}
if "priority" in tool_args and tool_args["priority"]:
properties["priority"] = {"name": tool_args["priority"]}
# Assignee: {"id": "uuid"}
if "assignee" in tool_args and tool_args["assignee"]:
properties["assignee"] = {"id": tool_args["assignee"]}
# Reporter: {"id": "uuid"}
if "reporter" in tool_args and tool_args["reporter"]:
properties["reporter"] = {"id": tool_args["reporter"]}
# Labels: [{"id": "uuid"}, ...]
if "labels" in tool_args and tool_args["labels"]:
labels = tool_args["labels"]
if isinstance(labels, list):
properties["labels"] = [{"id": label_id} for label_id in labels if label_id]
return properties
def map_state_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for state creation/update tools."""
properties: Dict[str, Any] = {}
# Color: {"name": "#FF5733"}
if "color" in tool_args and tool_args["color"]:
properties["color"] = {"name": tool_args["color"]}
# Group: {"name": "backlog/unstarted/started/completed/cancelled"}
if "group" in tool_args and tool_args["group"]:
properties["group"] = {"name": tool_args["group"]}
# Sequence: {"name": "1"}
if "sequence" in tool_args and tool_args["sequence"] is not None:
properties["sequence"] = {"name": str(tool_args["sequence"])}
# Is Triage: {"name": "true/false"}
if "is_triage" in tool_args and tool_args["is_triage"] is not None:
properties["is_triage"] = {"name": str(bool(tool_args["is_triage"])).lower()}
# Default: {"name": "true/false"}
if "default" in tool_args and tool_args["default"] is not None:
properties["default"] = {"name": str(bool(tool_args["default"])).lower()}
# External ID: {"name": "external-123"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
# External Source: {"name": "jira"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
# Description: {"name": "text"}
if "description" in tool_args and tool_args["description"]:
properties["description"] = {"name": tool_args["description"]}
return properties
def map_label_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for label creation/update tools."""
properties: Dict[str, Any] = {}
# Color: {"name": "#FF5733"}
if "color" in tool_args and tool_args["color"]:
properties["color"] = {"name": tool_args["color"]}
# Description: {"name": "text"}
if "description" in tool_args and tool_args["description"]:
properties["description"] = {"name": tool_args["description"]}
# Parent: {"id": "uuid"} - for nested labels
if "parent" in tool_args and tool_args["parent"]:
properties["parent"] = {"id": tool_args["parent"]}
# Sort Order: {"name": "1.5"}
if "sort_order" in tool_args and tool_args["sort_order"] is not None:
properties["sort_order"] = {"name": str(tool_args["sort_order"])}
# External ID: {"name": "external-123"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
# External Source: {"name": "jira"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
return properties
def map_worklog_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for worklog creation/update tools."""
properties: Dict[str, Any] = {}
# Description: {"name": "text"}
if "description" in tool_args and tool_args["description"]:
properties["description"] = {"name": tool_args["description"]}
# Duration: {"name": "60"} - duration in minutes
if "duration" in tool_args and tool_args["duration"] is not None:
properties["duration"] = {"name": str(tool_args["duration"])}
# Created By: {"id": "uuid"}
if "created_by" in tool_args and tool_args["created_by"]:
properties["created_by"] = {"id": tool_args["created_by"]}
# Updated By: {"id": "uuid"}
if "updated_by" in tool_args and tool_args["updated_by"]:
properties["updated_by"] = {"id": tool_args["updated_by"]}
return properties
def map_comment_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for comment creation/update tools."""
properties: Dict[str, Any] = {}
# Comment HTML: {"name": "<p>text</p>"}
if "comment_html" in tool_args and tool_args["comment_html"]:
properties["comment_html"] = {"name": tool_args["comment_html"]}
# Comment JSON: {"name": "{...}"} - converting dict to string representation
if "comment_json" in tool_args and tool_args["comment_json"] is not None:
import json
properties["comment_json"] = {"name": json.dumps(tool_args["comment_json"])}
# Access: {"name": "public/private/..."}
if "access" in tool_args and tool_args["access"]:
properties["access"] = {"name": tool_args["access"]}
# External ID: {"name": "external-123"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
# External Source: {"name": "jira"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
return properties
def map_link_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for link creation/update tools."""
properties: Dict[str, Any] = {}
# Title: {"name": "Link Title"}
if "title" in tool_args and tool_args["title"]:
properties["title"] = {"name": tool_args["title"]}
# URL: {"name": "https://example.com"}
if "url" in tool_args and tool_args["url"]:
properties["url"] = {"name": tool_args["url"]}
# Metadata: {"name": {...}} - JSON object
if "metadata" in tool_args and tool_args["metadata"]:
import json
if isinstance(tool_args["metadata"], dict):
properties["metadata"] = {"name": json.dumps(tool_args["metadata"])}
else:
properties["metadata"] = {"name": str(tool_args["metadata"])}
return properties
def map_attachment_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for attachment creation/update tools."""
properties: Dict[str, Any] = {}
# Name: {"name": "filename.pdf"}
if "name" in tool_args and tool_args["name"]:
properties["name"] = {"name": tool_args["name"]}
# Type: {"name": "application/pdf"}
if "type" in tool_args and tool_args["type"]:
properties["type"] = {"name": tool_args["type"]}
# Size: {"name": "1024"} - convert int to string
if "size" in tool_args and tool_args["size"] is not None:
properties["size"] = {"name": str(tool_args["size"])}
# External ID: {"name": "external-123"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
# External Source: {"name": "jira"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
# Is Uploaded: {"name": "true/false"} - for update operations
if "is_uploaded" in tool_args and tool_args["is_uploaded"] is not None:
properties["is_uploaded"] = {"name": str(bool(tool_args["is_uploaded"])).lower()}
return properties
def map_property_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for property creation/update tools."""
properties: Dict[str, Any] = {}
# Property Type: {"name": "TEXT/DATETIME/..."}
if "property_type" in tool_args and tool_args["property_type"]:
properties["property_type"] = {"name": tool_args["property_type"]}
# Relation Type: {"name": "ONE_TO_ONE/..."}
if "relation_type" in tool_args and tool_args["relation_type"]:
properties["relation_type"] = {"name": tool_args["relation_type"]}
# Is Required: boolean
if "is_required" in tool_args and tool_args["is_required"] is not None:
properties["is_required"] = bool(tool_args["is_required"])
# Default Value: list of strings (keep as-is)
if "default_value" in tool_args and tool_args["default_value"] is not None:
properties["default_value"] = tool_args["default_value"]
# Settings: dict/object (keep as-is)
if "settings" in tool_args and tool_args["settings"] is not None:
properties["settings"] = tool_args["settings"]
# Is Active: boolean
if "is_active" in tool_args and tool_args["is_active"] is not None:
properties["is_active"] = bool(tool_args["is_active"])
# Is Multi: boolean
if "is_multi" in tool_args and tool_args["is_multi"] is not None:
properties["is_multi"] = bool(tool_args["is_multi"])
# Validation Rules: dict/object (keep as-is)
if "validation_rules" in tool_args and tool_args["validation_rules"] is not None:
properties["validation_rules"] = tool_args["validation_rules"]
# External Source: {"name": "value"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
# External ID: {"name": "value"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
return properties
def map_type_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map properties for work item type creation/update tools."""
properties: Dict[str, Any] = {}
# Project IDs: [{"id": "uuid"}, ...]
if "project_ids" in tool_args and tool_args["project_ids"]:
project_ids = tool_args["project_ids"]
if isinstance(project_ids, list):
properties["project_ids"] = [{"id": project_id} for project_id in project_ids if project_id]
# Is Epic: boolean
if "is_epic" in tool_args and tool_args["is_epic"] is not None:
properties["is_epic"] = bool(tool_args["is_epic"])
# Is Active: boolean
if "is_active" in tool_args and tool_args["is_active"] is not None:
properties["is_active"] = bool(tool_args["is_active"])
# External Source: {"name": "value"}
if "external_source" in tool_args and tool_args["external_source"]:
properties["external_source"] = {"name": tool_args["external_source"]}
# External ID: {"name": "value"}
if "external_id" in tool_args and tool_args["external_id"]:
properties["external_id"] = {"name": tool_args["external_id"]}
return properties
def map_common_properties(tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""Map common properties that might appear in various tools."""
properties: Dict[str, Any] = {}
# Handle generic 'name' if it wasn't popped earlier or is needed as property
# Note: 'name' is usually a root parameter, but sometimes duplicated in properties
return properties
def map_tool_properties(artifact_type: str, tool_args: Dict[str, Any]) -> Dict[str, Any]:
"""
Main entry point to map tool arguments to the nested 'properties' structure.
Args:
artifact_type: Type of the artifact being executed
tool_args: Dictionary of arguments passed to the tool
Returns:
Dictionary representing the 'properties' field in action summary
"""
properties: Dict[str, Any] = {}
# Dispatch based on tool name pattern
if artifact_type == "workitem" or artifact_type == "epic":
properties.update(map_workitem_properties(tool_args))
elif artifact_type == "project":
properties.update(map_project_properties(tool_args))
elif artifact_type == "module":
properties.update(map_module_properties(tool_args))
elif artifact_type == "cycle":
properties.update(map_cycle_properties(tool_args))
elif artifact_type == "intake":
properties.update(map_intake_properties(tool_args))
elif artifact_type == "label":
properties.update(map_label_properties(tool_args))
elif artifact_type == "page":
return {} # Page doesn't have any properties other than the basic ones already added in params dict
elif artifact_type == "state":
properties.update(map_state_properties(tool_args))
elif artifact_type == "worklog":
properties.update(map_worklog_properties(tool_args))
elif artifact_type == "comment":
properties.update(map_comment_properties(tool_args))
elif artifact_type == "attachment":
properties.update(map_attachment_properties(tool_args))
elif artifact_type == "link":
properties.update(map_link_properties(tool_args))
elif artifact_type == "type":
properties.update(map_type_properties(tool_args))
elif artifact_type == "property":
properties.update(map_property_properties(tool_args))
# Add any other specific mappers here
# Fallback/Generic handling for remaining keys could be added here
# but for now we stick to explicit mappings to avoid cluttering the UI
return properties
apps/pi/pi/services/chat/helpers/action_summary_generator.py
-1417
View File
File diff suppressed because it is too large Load Diff
apps/pi/pi/services/chat/helpers/agent_executor.py
-791
View File
@@ -1,791 +0,0 @@
"""Agent execution logic for single and multi-agent scenarios."""
import ast
import asyncio
import re
from collections.abc import AsyncIterator
from typing import Any
from typing import Dict
from typing import List
from typing import Tuple
from typing import Union
from typing import cast
from langchain_core.messages import AIMessage
from langchain_core.messages import BaseMessage
from langchain_core.messages import HumanMessage
from langchain_core.messages import SystemMessage
from langchain_core.messages import ToolMessage
from pi import logger
from pi.app.models.enums import ExecutionStatus
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import MessageMetaStepType
from pi.core.db.plane_pi.lifecycle import get_streaming_db_session
from pi.services.chat.prompts import multi_tool_system_prompt
from pi.services.chat.utils import StandardAgentResponse
from pi.services.chat.utils import format_conversation_history
from pi.services.chat.utils import get_current_timestamp_context
from pi.services.chat.utils import mask_uuids_in_text
from pi.services.chat.utils import standardize_flow_step_content
from pi.services.retrievers.pg_store.json_serializer import sanitize_execution_data
from pi.services.retrievers.pg_store.message import upsert_message_flow_steps as _upsert_message_flow_steps
from pi.services.schemas.chat import Agents
log = logger.getChild(__name__)
async def execute_single_agent(
chatbot_instance,
agent,
sub_query,
user_meta,
message_id,
workspace_id,
project_id,
conversation_history,
user_id,
chat_id,
query_flow_store,
parsed_query,
query_id,
step_order,
workspace_in_context,
db,
preset_tables=None,
preset_sql_query=None,
preset_placeholders=None,
):
"""Execute a single agent and prepare response stream."""
if not workspace_in_context:
return chatbot_instance.handle_generic_query_stream(
sub_query,
user_id,
conversation_history,
user_meta,
non_plane=True,
include_user_context_on_followup=bool(conversation_history),
), None
if agent == Agents.GENERIC_AGENT:
return chatbot_instance.handle_generic_query_stream(sub_query, user_id, conversation_history, user_meta), None
agent_result = await chatbot_instance.handle_agent_query(
db,
agent,
sub_query,
user_meta,
message_id,
workspace_id,
project_id,
conversation_history,
user_id,
chat_id,
query_flow_store,
is_multi_agent=False,
preset_tables=preset_tables,
preset_sql_query=preset_sql_query,
preset_placeholders=preset_placeholders,
)
# Define response variable with proper typing upfront
response: Union[str, Dict[str, Any]]
if isinstance(agent_result, tuple):
intermediate_results, response = agent_result
else:
intermediate_results = None
response = agent_result
responses: List[Tuple[str, str, Union[str, Dict[str, Any]]]] = [(agent, sub_query, response)]
formatted_responses_result = StandardAgentResponse.format_responses(responses, chat_id, query_flow_store)
# Check if responses contain errors - if so, return error directly without LLM rewriting
if isinstance(formatted_responses_result, dict) and formatted_responses_result.get("has_errors"):
# Extract clean error message from the response
error_message = StandardAgentResponse.extract_results(response)
if not error_message:
error_message = "An error occurred while retrieving data. Please try again later."
# Clean up any markdown formatting from SQL agent
error_message = error_message.strip()
if error_message.startswith("- **Result**: "):
error_message = error_message.replace("- **Result**: ", "")
return chatbot_instance._create_simple_stream(error_message), None
# Extract formatted string for normal flow
formatted_responses = formatted_responses_result.get("formatted") if isinstance(formatted_responses_result, dict) else formatted_responses_result
formatted_history = format_conversation_history(conversation_history)
# Sanitize execution_data to ensure JSON serializability
sanitized_execution_data = sanitize_execution_data(intermediate_results or {})
flow_steps = [
{
"step_order": step_order,
"step_type": FlowStepType.TOOL.value,
"tool_name": agent.name,
"content": standardize_flow_step_content(response, FlowStepType.TOOL),
"execution_data": sanitized_execution_data,
}
]
try:
async with get_streaming_db_session() as _subdb:
flow_step_result = await asyncio.wait_for(
_upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=flow_steps,
db=_subdb,
),
timeout=2.0,
)
except asyncio.TimeoutError:
log.warning(f"Timed out recording flow step for message {query_id}; continuing")
flow_step_result = {"message": "error", "error": "timeout"}
except Exception as e:
log.warning(f"Failed to record flow step for message {query_id}: {e}")
flow_step_result = {"message": "error", "error": str(e)}
if flow_step_result["message"] != "success":
# Log the error but continue - flow step logging is not critical for user experience
log.warning(f"Failed to record flow step for message {query_id}: {flow_step_result.get("error", "Unknown error")}")
return chatbot_instance.combined_response_stream(
parsed_query,
formatted_responses,
formatted_history,
user_meta,
user_id,
workspace_in_context=workspace_in_context,
), None
async def execute_multi_agents(
chatbot_instance,
selected_agents,
user_meta,
workspace_id,
workspace_slug,
project_id,
conversation_history,
enhanced_conversation_history,
user_id,
chat_id,
query_flow_store,
parsed_query,
query_id,
step_order,
db,
original_query,
reasoning_container=None,
workspace_in_context: bool | None = None,
) -> AsyncIterator[str]:
"""Execute multiple agents using LangChain tool calling for orchestration with real-time streaming."""
log.info(f"ChatID: {chat_id} - Starting multi-agent execution with {len(selected_agents)} agent(s): {[agent.agent for agent in selected_agents]}")
# Reset shared state for this query
chatbot_instance.vector_search_issue_ids = []
chatbot_instance.vector_search_page_ids = []
# Reset stored agent responses for clean state
chatbot_instance.agent_responses = {}
# Get tools for selected agents
log.info(f"ChatID: {chat_id} - Fetching tools for selected agents")
tools = await chatbot_instance._get_selected_tools(
db,
selected_agents,
user_meta,
workspace_id,
workspace_slug,
project_id,
user_id,
chat_id,
query_flow_store,
conversation_history,
query_id,
)
if not tools:
log.error(f"ChatID: {chat_id} - No valid tools found for selected agents")
# Finalize with a user-friendly message so caller can persist
yield "πspecial reasoning blockπ: 📝 Generating final response...\n\n"
yield "An unexpected error occurred. Please try again later."
return
custom_prompt = (
"⚠️ IMPORTANT: The router has suggested the following tools, but you MUST override "
"this selection if a more specialized tool is available.\n\n"
"Suggested tools and their queries: \n"
)
# Map agent names to actual tool names to match bound tools
agent_to_tool_map = {
"plane_structured_database_agent": "structured_db_tool",
"plane_vector_search_agent": "vector_search_tool",
"plane_pages_agent": "pages_search_tool",
"plane_docs_agent": "docs_search_tool",
"generic_agent": "generic_query_tool",
}
tool_entries = []
for agent in selected_agents:
tool_name = agent_to_tool_map.get(agent.agent, agent.agent)
tool_entries.append(f"Tool Name: {tool_name};\nTool Query: {agent.query}")
custom_prompt += "\n\n".join(tool_entries)
# Inject enhanced conversation history if available to guide orchestration
try:
if enhanced_conversation_history and isinstance(enhanced_conversation_history, str) and enhanced_conversation_history.strip():
custom_prompt += f"\n\n**CONVERSATION HISTORY & ACTION CONTEXT:**\n{enhanced_conversation_history}\n"
except Exception:
pass
# Provide PROJECT/USER/TIME context to the tool planner
context_block = ""
try:
# Project scoping (project chat)
if project_id:
context_block += f"\n\n**🔥 PROJECT CONTEXT (CRITICAL):**\nProject ID: {project_id}\n\n**IMPORTANT SCOPING RULES:**\n- This is a PROJECT-LEVEL chat - ALL operations are scoped to THIS PROJECT ONLY\n- When the request mentions 'current cycle', 'current module', 'work items', etc. - it means ONLY within THIS PROJECT\n- Use this project_id for ALL tools that accept project_id parameter\n- DO NOT query across all projects - scope everything to THIS specific project\n- User refers to 'this project'/'the project'/'current project' = use this project_id" # noqa: E501
else:
# Workspace-level context (no specific project)
context_block += f"\n\n**🌐 WORKSPACE CONTEXT (CRITICAL):**\nWorkspace ID: {workspace_id}\n\n**IMPORTANT SCOPING RULES:**\n- This is a WORKSPACE-LEVEL chat - queries can span MULTIPLE PROJECTS\n- When the request mentions 'last cycle', 'this cycle', 'work items', etc. WITHOUT specifying a project - it could be in ANY project\n- Use list_member_projects (without limit or with high limit) to get ALL projects in the workspace\n- Then iterate through projects to find relevant cycles/modules/work-items\n- CRITICAL: Do NOT limit to 1 project unless the user specifically names or refers to a specific project" # noqa: E501
# User context
if user_id:
context_block += "\n**USER CONTEXT:**\n" f"User ID: {user_id}\n" "- Use this to resolve pronouns like 'me', 'my', 'assigned to me'.\n"
# Date/time context
dt_ctx = None
try:
dt_ctx = await get_current_timestamp_context(user_id)
except Exception:
dt_ctx = None
if dt_ctx:
context_block += f"\n{dt_ctx}\n"
except Exception:
# Non-fatal; continue without context block if any error occurs
pass
# CRITICAL: Inject clarification context if this is a follow-up to ask_for_clarification
if user_meta and isinstance(user_meta, dict) and user_meta.get("clarification_context"):
clar_ctx = user_meta["clarification_context"]
original_query_text = clar_ctx.get("original_query")
reason = clar_ctx.get("reason")
disambig_options = clar_ctx.get("disambiguation_options") or []
answer_text = clar_ctx.get("answer_text")
custom_prompt += "\n\n**CLARIFICATION CONTEXT:**\n"
# CRITICAL: Include the original query to maintain full context
if original_query_text:
custom_prompt += f"Original user request: {original_query_text}\n"
if reason:
custom_prompt += f"Clarification reason: {reason}\n"
if disambig_options:
custom_prompt += "The user was previously shown these options:\n"
for idx, opt in enumerate(disambig_options, 1):
if isinstance(opt, dict):
opt_id = opt.get("id")
opt_name = opt.get("name") or opt.get("display_name") or ""
opt_identifier = opt.get("identifier") or ""
opt_email = opt.get("email") or ""
# Format based on what fields are available
if opt_email:
# User entity
custom_prompt += f" {idx}. {opt_name} ({opt_email}) → UUID: {opt_id}\n"
elif opt_identifier:
# Project/workitem entity
custom_prompt += f" {idx}. {opt_name} (Identifier: {opt_identifier}) → UUID: {opt_id}\n"
else:
# Generic entity
custom_prompt += f" {idx}. {opt_name} → UUID: {opt_id}\n"
if answer_text:
custom_prompt += f"\nUser's clarification answer: {answer_text}\n"
custom_prompt += "\nIMPORTANT: The current user message is a clarification response to the original request above. Use the clarification answer to resolve missing information and continue with the ORIGINAL request, not as a new standalone request.\n" # noqa E501
# Bind tools to the LLM
llm_with_tools = chatbot_instance.tool_llm.bind_tools(tools)
# Set tracking context on the bound LLM instance
llm_with_tools.set_tracking_context(query_id, db, MessageMetaStepType.TOOL_ORCHESTRATION, chat_id=str(chat_id))
# Track execution
tool_flow_steps: List[Dict[str, Any]] = []
current_step = step_order
messages: List[BaseMessage] = []
responses: List[Tuple[str, str, Union[str, Dict[str, Any]]]] = []
messages.append(SystemMessage(content=multi_tool_system_prompt))
# Prepend context at the top of the prompt for better LLM comprehension
if context_block.strip():
custom_prompt = f"{context_block}\n\n{custom_prompt}"
try:
query_to_use = parsed_query
messages.append(HumanMessage(content=custom_prompt + f"\n\nUser Query: {query_to_use}"))
# Record the tool orchestration context (enhanced conversation history) before any tool planning/execution
try:
if enhanced_conversation_history and isinstance(enhanced_conversation_history, str) and enhanced_conversation_history.strip():
tool_flow_steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": "tool_orchestration_context",
"content": "Context used for tool orchestration",
"execution_data": {"enhanced_conversation_history": enhanced_conversation_history},
"is_planned": False,
"is_executed": False,
})
current_step += 1
except Exception:
pass
# log.info("************************************************")
# log.info(f"ChatID: {chat_id} - multi-agent custom prompt: {custom_prompt}")
# log.info("************************************************")
# Log what's sent to the tool selection LLM
log.info(f"ChatID: {chat_id} - Tool selection LLM input - Query: {query_to_use}")
log.info(f"ChatID: {chat_id} - Tool selection LLM input - Available tools: {[t.name for t in tools]}")
# Yield initial status
reasoning_chunk = "🤖 Retrieving information...\n\n"
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk
yield f"πspecial reasoning blockπ: {reasoning_chunk}"
# Initial invocation of LLM with tools
llm_response = await chatbot_instance._tool_orchestration_llm_call(llm_with_tools, messages)
# Handle failure case
if llm_response == "TOOL_ORCHESTRATION_FAILURE":
# Persist any steps captured so far
if tool_flow_steps:
try:
async with get_streaming_db_session() as _subdb:
await _upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=tool_flow_steps,
db=_subdb,
)
except Exception as _e:
log.warning(f"ChatID: {chat_id} - Failed to persist tool flow steps on failure: {_e}")
# Finalize with a short fallback message so caller can persist
yield "πspecial reasoning blockπ: 📝 Generating final response...\n\n"
yield "An unexpected error occurred. Please try again later."
return
ai_message = cast(AIMessage, llm_response)
messages.append(ai_message)
if not ai_message.tool_calls:
log.info(f"ChatID: {chat_id} - No tool calls found in the initial LLM response. {ai_message.content}")
responses.append((Agents.GENERIC_AGENT, query_to_use, ai_message.content)) # type: ignore
else:
# Log the tool calls selected by the LLM
tool_calls_summary = [(tc["name"], tc.get("args", {})) for tc in ai_message.tool_calls]
log.info(f"ChatID: {chat_id} - Tool selection LLM chose {len(ai_message.tool_calls)} tool(s): {tool_calls_summary}")
# Process tool calls iteratively
while ai_message.tool_calls:
# Execute all tool calls
tool_messages = []
for tool_call in ai_message.tool_calls:
tool_name = tool_call["name"]
tool_args = tool_call["args"]
tool_id = tool_call["id"]
# Intercept clarification requests and short-circuit
if tool_name == "ask_for_clarification":
log.info(f"ChatID: {chat_id} - Clarification requested by LLM: {tool_args.get("reason", "No reason provided")}")
try:
# Execute the clarification tool to get payload
tool_to_execute = next((t for t in tools if t.name == tool_name), None)
result = await tool_to_execute.ainvoke(tool_args) if tool_to_execute else "{}"
except Exception as e:
log.error(f"ChatID: {chat_id} - Clarification tool failed: {e}")
result = "{}"
# Persist flow step with clarification_pending
try:
from pi.services.chat.utils import standardize_flow_step_content as _std
payload = str(result)
clarification_execution_data = {
"args": tool_args,
"clarification_pending": True,
"clarification_payload": payload,
# Store planning context to avoid re-routing on follow-up
"selected_categories": [agent.agent for agent in selected_agents],
"bound_tool_names": [t.name for t in tools],
# CRITICAL: Store the original query so we can reconstruct full context on clarification follow-up
"original_query": original_query,
}
# Sanitize execution_data to ensure JSON serializability
sanitized_clarification_data = sanitize_execution_data(clarification_execution_data)
tool_flow_steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL,
"tool_name": "ask_for_clarification",
"content": _std(payload, FlowStepType.TOOL),
"execution_data": sanitized_clarification_data,
"is_planned": False,
"is_executed": False,
"execution_success": ExecutionStatus.PENDING,
})
current_step += 1
except Exception:
pass
# Persist immediately since we return right after streaming
try:
async with get_streaming_db_session() as _subdb:
flow_steps_result = await _upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=tool_flow_steps,
db=_subdb,
)
if flow_steps_result.get("message") != "success":
log.warning("Failed to record clarification flow step")
except Exception as e:
log.error(f"ChatID: {chat_id} - Failed to persist clarification flow step: {e}")
# Stream clarification block and stop orchestration
try:
yield f"πspecial clarification blockπ: {result}\n"
except Exception:
yield f"πspecial clarification blockπ: {str(result)}\n"
return
# Yield tool execution start (robust to different arg schemas)
if tool_name != "generic_query_tool":
try:
from pi.services.chat.helpers.tool_utils import format_tool_query_for_display
user_friendly_tool_name = chatbot_instance._tool_name_shown_to_user(tool_name)
# Build a readable query snippet depending on tool schema
display_arg = format_tool_query_for_display(tool_name, tool_args, user_query=query_to_use)
masked_display = mask_uuids_in_text(display_arg or (query_to_use or "the request"))
reasoning_chunk = f"🔧 Performing {user_friendly_tool_name.lower()} to pull details about: {masked_display}\n\n"
except Exception:
user_friendly_tool_name = chatbot_instance._tool_name_shown_to_user(tool_name)
reasoning_chunk = f"🔧 Performing {user_friendly_tool_name.lower()}\n\n"
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk
yield f"πspecial reasoning blockπ: {reasoning_chunk}"
# Find and execute the tool
tool_to_execute = next((t for t in tools if t.name == tool_name), None)
if tool_to_execute:
log.info(f"ChatID: {chat_id} - Executing tool: {tool_name} with args: {str(tool_args)[:100]}")
try:
# Execute the tool directly; endpoint-level SSE heartbeat will keep the stream alive
tool_result = await tool_to_execute.ainvoke(tool_args)
# Log brief output for observability
result_preview = str(tool_result)[:200] if tool_result else "None"
log.info(f"ChatID: {chat_id} - Tool {tool_name} result preview: {result_preview}")
log.debug(f"ChatID: {chat_id} - Tool {tool_name} completed successfully")
# Yield tool completion and result
if not tool_name == "generic_query_tool":
user_friendly_tool_name = chatbot_instance._tool_name_shown_to_user(tool_name)
reasoning_chunk = f"✅ {user_friendly_tool_name} execution completed\n\n"
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk
yield f"πspecial reasoning blockπ: {reasoning_chunk}"
# yield f"📊 Result: {str(tool_result)[:200]}{'...' if len(str(tool_result)) > 200 else ''}\n\n"
# Record the tool execution with standardized execution_data format
execution_data: Dict[str, Any] = {
"tool_args": tool_args,
"tool_name": tool_name,
"execution_status": "success",
# "timestamp": await get_current_timestamp_context(user_id),
}
stored_response = chatbot_instance.agent_responses.get(tool_name)
if tool_name == "structured_db_tool":
# For structured_db_tool, include full intermediate_results for SQL debugging
if stored_response and "intermediate_results" in stored_response:
execution_data["intermediate_results"] = stored_response["intermediate_results"]
# Extract key debugging info to top level for easier access
ir = stored_response["intermediate_results"]
execution_data["sql_query"] = ir.get("generated_sql", "")
execution_data["final_query"] = ir.get("final_query", "")
execution_data["relevant_tables"] = ir.get("relevant_tables", [])
execution_data["tool_type"] = "database_query"
elif tool_name in ["vector_search_tool", "pages_search_tool", "docs_search_tool"]:
# For search tools, standardize search metadata
if stored_response:
entity_urls = stored_response.get("entity_urls") or []
execution_data["entity_urls"] = entity_urls
execution_data["results_count"] = len(entity_urls)
execution_data["search_query"] = tool_args.get("query", "")
# Include execution metadata if available
if "execution_metadata" in stored_response:
execution_data["search_metadata"] = stored_response["execution_metadata"]
execution_data["tool_type"] = "semantic_search"
elif tool_name == "generic_query_tool":
execution_data["tool_type"] = "generic_llm"
if stored_response and isinstance(stored_response, dict):
execution_data.update({k: v for k, v in stored_response.items() if k not in ["results", "tool_args"]})
else:
# For other tools, include any available stored response data
execution_data["tool_type"] = "unknown"
if stored_response and isinstance(stored_response, dict):
execution_data.update({k: v for k, v in stored_response.items() if k not in ["results", "tool_args"]})
# Sanitize execution_data to ensure JSON serializability
sanitized_execution_data = sanitize_execution_data(execution_data)
# Extract key facts from tool_result and enrich execution_data
try:
facts: Dict[str, Any] = {}
result_str = str(tool_result) if tool_result is not None else ""
# Extract explicit entity lines from success formatter
m_url = re.search(r"^Entity URL:\s*(.+)$", result_str, flags=re.MULTILINE)
if m_url:
facts["entity_url"] = m_url.group(1).strip()
m_eid = re.search(r"^Entity ID:\s*(.+)$", result_str, flags=re.MULTILINE)
if m_eid:
facts["entity_id"] = m_eid.group(1).strip()
m_ename = re.search(r"^Entity Name:\s*(.+)$", result_str, flags=re.MULTILINE)
if m_ename:
facts["entity_name"] = m_ename.group(1).strip()
m_etype = re.search(r"^Entity Type:\s*(.+)$", result_str, flags=re.MULTILINE)
if m_etype:
facts["entity_type"] = m_etype.group(1).strip()
m_eident = re.search(r"^Entity Identifier:\s*(.+)$", result_str, flags=re.MULTILINE)
if m_eident:
facts["entity_identifier"] = m_eident.group(1).strip()
# Extract Result: payload and try to parse dict/list safely
m_res = re.search(r"\n\nResult:\s*(\{[\s\S]*?\}|\[[\s\S]*?\])", result_str)
if m_res:
payload = m_res.group(1)
try:
parsed = ast.literal_eval(payload)
if isinstance(parsed, dict):
for k in ("id", "identifier", "project_id", "workspace_id", "name"):
v = parsed.get(k)
if v:
facts[k] = v
# Common list containers
for list_key in ("projects", "cycles", "results", "items"):
if isinstance(parsed.get(list_key), list):
lst = parsed.get(list_key) or []
facts[f"{list_key}_count"] = len(lst)
ids = []
for it in lst[:3]:
if isinstance(it, dict) and it.get("id"):
ids.append(str(it.get("id")))
if ids:
facts[f"{list_key}_ids_preview"] = ids
elif isinstance(parsed, list):
facts["list_count"] = len(parsed)
ids = []
for it in parsed[:3]:
if isinstance(it, dict) and it.get("id"):
ids.append(str(it.get("id")))
if ids:
facts["ids_preview"] = ids
except Exception:
pass
# Always include a concise args summary as fallback
try:
if isinstance(tool_args, dict):
arg_keys = [
k for k in ("project_id", "module_id", "cycle_id", "issue_id", "name", "identifier") if k in tool_args
]
if arg_keys:
facts["args_summary"] = {k: tool_args.get(k) for k in arg_keys}
except Exception:
pass
if facts:
# Merge into sanitized_execution_data
if isinstance(sanitized_execution_data, dict):
sanitized_execution_data["facts"] = facts
except Exception:
pass
tool_flow_steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": tool_name,
"content": standardize_flow_step_content(tool_result, FlowStepType.TOOL),
"execution_data": sanitized_execution_data,
"execution_success": ExecutionStatus.SUCCESS, # Mark as SUCCESS after successful execution
})
current_step += 1
# Add to responses for final formatting
agent_name = chatbot_instance._tool_name_to_agent(tool_name)
# Generic approach: Check if any agent has stored entity URLs
response_with_urls: Union[str, Dict[str, Any]] = tool_result
stored_response = chatbot_instance.agent_responses.get(tool_name)
if stored_response and stored_response.get("entity_urls"):
response_with_urls = stored_response
responses.append((agent_name, str(tool_args), response_with_urls))
# Create tool message
tool_message = ToolMessage(content=str(tool_result), tool_call_id=tool_id or "")
tool_messages.append(tool_message)
except Exception as e:
log.error(f"ChatID: {chat_id} - Error executing tool {tool_name}: {str(e)}")
# Use a generic, user-friendly reasoning message
reasoning_chunk = "❌ Encountered an unexpected error.\n\n"
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk
yield f"πspecial reasoning blockπ: {reasoning_chunk}"
# Record the failed tool execution with error details
error_execution_data = {
"tool_args": tool_args,
"tool_name": tool_name,
"execution_status": "error",
"error_message": str(e),
"tool_type": "unknown",
}
# Sanitize error execution_data to ensure JSON serializability
sanitized_error_execution_data = sanitize_execution_data(error_execution_data)
tool_flow_steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL,
"tool_name": tool_name,
"content": f"Error: {str(e)}",
"execution_data": sanitized_error_execution_data,
"execution_success": ExecutionStatus.FAILED, # Mark as FAILED on error
"execution_error": str(e),
})
current_step += 1
error_message = ToolMessage(content=f"Error executing {tool_name}: {str(e)}", tool_call_id=str(tool_id or ""))
tool_messages.append(error_message)
else:
log.error(f"ChatID: {chat_id} - Tool {tool_name} not found")
# Generic, user-friendly status instead of internal details
reasoning_chunk = "⚠️ A required capability was unavailable.\n\n"
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk
yield f"πspecial reasoning blockπ: {reasoning_chunk}"
error_message = ToolMessage(content=f"Tool {tool_name} not found", tool_call_id=str(tool_id or ""))
tool_messages.append(error_message)
# Add tool messages to conversation
messages.extend(tool_messages)
# Yield status before getting next LLM response
reasoning_chunk = "🤖 Analyzing results...\n\n"
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk
yield f"πspecial reasoning blockπ: {reasoning_chunk}"
# Get next response from LLM
llm_response = await chatbot_instance._tool_orchestration_llm_call(llm_with_tools, messages)
# Handle failure case
if llm_response == "TOOL_ORCHESTRATION_FAILURE":
# Finalize with a short fallback message so caller can persist
yield "πspecial reasoning blockπ: 📝 Generating final response...\n\n"
yield "An unexpected error occurred. Please try again later."
return
ai_message = cast(AIMessage, llm_response)
messages.append(ai_message)
# Log any follow-up tool calls
if ai_message.tool_calls:
tool_calls_summary = [(tc["name"], tc.get("args", {})) for tc in ai_message.tool_calls]
log.info(f"ChatID: {chat_id} - Tool selection LLM follow-up: {len(ai_message.tool_calls)} tool(s): {tool_calls_summary}")
# Add all tool flow steps to database
if tool_flow_steps:
try:
async with get_streaming_db_session() as _subdb:
flow_steps_result = await asyncio.wait_for(
_upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=tool_flow_steps,
db=_subdb,
),
timeout=2.0,
)
except asyncio.TimeoutError:
log.warning(f"Timed out recording tool flow steps for message {query_id}; continuing")
flow_steps_result = {"message": "error", "error": "timeout"}
except Exception as e:
log.warning(f"Failed to record tool flow steps for message {query_id}: {e}")
flow_steps_result = {"message": "error", "error": str(e)}
if flow_steps_result["message"] != "success":
# Log the error but continue - flow step logging is not critical for user experience
log.warning(f"Failed to record flow steps for message {query_id}: {flow_steps_result.get("error", "Unknown error")}")
# Continue with response generation - user should still get their answer
# Yield final status
reasoning_chunk = "📝 Generating final response...\n\n"
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk
yield f"πspecial reasoning blockπ: {reasoning_chunk}"
# Format responses for the combination prompt
log.info(f"ChatID: {chat_id} - Multi-agent execution completed. Processing {len(responses)} response(s) from agents")
formatted_responses_result = StandardAgentResponse.format_responses(responses, chat_id, query_flow_store)
# Check if responses contain errors - if so, return error directly without LLM rewriting
if isinstance(formatted_responses_result, dict) and formatted_responses_result.get("has_errors"):
log.warning(f"ChatID: {chat_id} - Error detected in agent responses, bypassing LLM combination")
# Use friendly generic message for end users
clean_error_msg: str = "An unexpected error occurred. Please try again later."
# Persist steps before returning early
if tool_flow_steps:
try:
async with get_streaming_db_session() as _subdb:
await _upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=tool_flow_steps,
db=_subdb,
)
except Exception as _e:
log.warning(f"ChatID: {chat_id} - Failed to persist tool flow steps on retrieval error: {_e}")
# Finalize so caller can persist a proper assistant message
yield "πspecial reasoning blockπ: 📝 Generating final response...\n\n"
yield clean_error_msg
return
# Extract formatted string for normal flow
log.info(f"ChatID: {chat_id} - All agents completed successfully, proceeding to LLM combination")
formatted_responses = (
formatted_responses_result.get("formatted") if isinstance(formatted_responses_result, dict) else formatted_responses_result
)
formatted_history = format_conversation_history(conversation_history)
# Stream the combined response
log.debug(f"ChatID: {chat_id} - Starting combined_response_stream with {len(conversation_history)} history messages")
async for chunk in chatbot_instance.combined_response_stream(
parsed_query,
formatted_responses,
formatted_history,
user_meta,
user_id,
workspace_in_context=workspace_in_context,
):
yield chunk
log.info(f"ChatID: {chat_id} - Multi-agent flow completed successfully")
except Exception as e:
log.error(f"ChatID: {chat_id} - Error in LangChain tool orchestration: {str(e)}")
yield "πspecial reasoning blockπ: 📝 Generating final response...\n\n"
yield "An unexpected error occurred. Please try again later."
apps/pi/pi/services/chat/helpers/agent_selector.py
-187
View File
@@ -1,187 +0,0 @@
"""Agent selection and routing logic."""
import time
from typing import List
from typing import Tuple
from typing import Type
from typing import Union
from typing import cast
from langchain_core.prompts import ChatPromptTemplate
from openai import LengthFinishReasonError # type: ignore[attr-defined]
from pydantic.v1 import BaseModel
from pi import logger
from pi import settings
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import MessageMetaStepType
from pi.services.chat.prompts import router_prompt
from pi.services.chat.utils import standardize_flow_step_content
from pi.services.llm.llms import LLMConfig
from pi.services.llm.llms import create_openai_llm
from pi.services.retrievers.pg_store.message import upsert_message_flow_steps
from pi.services.schemas.chat import AgentQuery
from pi.services.schemas.chat import Agents
from pi.services.schemas.chat import RouteQuery
from pi.services.schemas.chat import RoutingResult
log = logger.getChild(__name__)
async def select_agents(
chatbot_instance, target, parsed_query, enhanced_conversation_history, workspace_in_context, query_id, chat_id, step_order, db
) -> Tuple[List[AgentQuery], int, Union[str, None]]:
"""Select appropriate agents based on query and context."""
selected_agents: list[AgentQuery] = []
error_message = None
if workspace_in_context:
# Prepare context-aware prompt for the enhanced router
if enhanced_conversation_history:
custom_prompt = f"Here is the conversation history for context:\n\n{enhanced_conversation_history}\n\nNew user query: {parsed_query}"
else:
custom_prompt = f"Here is the user query: {parsed_query}"
if target:
# Append note about resolved entity ID to existing prompt
note = f"""Please note that the user has mentioned a specific entity of type {target} using the 'mentions'(@entity_name) syntax in their query.
I resolved the UUID of this entity and enhanced the user query with it above.
Since this is a UUID, please always include the structured DB agent in the selected agents list.""" # noqa: E501
custom_prompt = f"{custom_prompt}\n\n{note}"
log.info("************************************************")
log.info(f"ChatID: {chat_id} - Router custom prompt: {custom_prompt}")
log.info("************************************************")
# Create router dynamically with tracking context
route_query_type: Type[BaseModel] = RouteQuery # type: ignore[assignment]
structured_decomposer = chatbot_instance.decomposer_llm.with_structured_output(
route_query_type,
include_raw=True,
method="json_mode",
)
structured_decomposer.set_tracking_context(query_id, db, MessageMetaStepType.ROUTER, chat_id=str(chat_id))
# Create dynamic router chain
json_schema_hint = '{"decomposed_queries": [{"agent": "<agent_name>", "query": "<query_text>"}]}'
router_prompt_template = ChatPromptTemplate.from_messages([
(
"system",
"{base_prompt}\n\nRespond **only** in valid JSON format matching this schema:\n{json_schema_hint}",
),
("human", "{custom_prompt}"),
]).partial(base_prompt=router_prompt, json_schema_hint=json_schema_hint)
dynamic_router = router_prompt_template | structured_decomposer
# handle openai.LengthFinishReasonError by retrying once with a slightly higher temperature
try:
router_start = time.time()
log.info(f"ChatID: {chat_id} - Starting router LLM call")
routing_result = cast(RoutingResult, await dynamic_router.ainvoke({"custom_prompt": custom_prompt}))
router_elapsed = time.time() - router_start
log.info(f"ChatID: {chat_id} - Router LLM call completed in {router_elapsed:.2f}s")
except LengthFinishReasonError as e:
log.error(f"Error during routing: {e}")
# Determine a new temperature (current + 0.1, capped at 1.0)
base_temp: float = 0.0 # decomposer_llm default temperature
new_temp: float = min(round(base_temp + 0.1, 2), 1.0)
log.debug("Retrying routing with increased temperature. Old: %s, New: %s", base_temp, new_temp)
# Create new decomposer LLM with higher temperature and structured output
# Switch to GPT-4o to mitigate length-finish errors during structured parsing
retry_decomposer_config = LLMConfig(model=settings.llm_model.GPT_4O, temperature=new_temp, streaming=False)
retry_decomposer_llm = create_openai_llm(retry_decomposer_config)
# Create new structured output and router
# Use the same type alias; mypy: no-redef warning avoided by reusing existing name
route_query_type_retry: Type[BaseModel] = RouteQuery # type: ignore[assignment]
retry_json_llm = retry_decomposer_llm.with_structured_output(
route_query_type_retry,
include_raw=True,
method="json_mode",
)
# Set tracking context on the structured LLM instance
retry_json_llm.set_tracking_context(query_id, db, MessageMetaStepType.ROUTER, chat_id=str(chat_id))
retry_prompt = ChatPromptTemplate.from_messages([
("system", router_prompt),
("human", "{custom_prompt}"),
])
retry_router = retry_prompt | retry_json_llm
retry_start = time.time()
log.info(f"ChatID: {chat_id} - Starting retry router LLM call")
routing_result = cast(RoutingResult, await retry_router.ainvoke({"custom_prompt": custom_prompt}))
retry_elapsed = time.time() - retry_start
log.info(f"ChatID: {chat_id} - Retry router LLM call completed in {retry_elapsed:.2f}s")
if routing_result["parsed"] is not None:
parsed_result = routing_result["parsed"]
selected_agents = parsed_result.decomposed_queries
log.info(f"ChatID: {chat_id} - Selected agents before target check: {selected_agents}")
if target:
# look if the actions agent is not in the selected_agents list and plane_structured_database_agent is in the selected_agents list
if Agents.PLANE_ACTION_EXECUTOR_AGENT not in [
agent.agent.name for agent in selected_agents
] and Agents.PLANE_STRUCTURED_DATABASE_AGENT in [agent.agent.name for agent in selected_agents]:
selected_agents.append(AgentQuery(agent=Agents.PLANE_STRUCTURED_DATABASE_AGENT, query=parsed_query))
log.info(f"ChatID: {chat_id} - Added plane_structured_database_agent to selected agents: {selected_agents}")
routing_content = [{"tool": agent.agent.name, "query": agent.query} for agent in selected_agents]
# Log enhanced history storage
history_preview = enhanced_conversation_history[:200] if enhanced_conversation_history else "(empty)"
log.info(
f"ChatID: {chat_id} - Storing enhanced conversation history in routing step (length={len(enhanced_conversation_history or "")}): {history_preview}" # noqa: E501
)
flow_step_result = await upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=[
{
"step_order": step_order,
"step_type": FlowStepType.ROUTING.value,
"tool_name": None,
"content": standardize_flow_step_content(routing_content, FlowStepType.ROUTING),
"execution_data": {"enhanced_conversation_history": enhanced_conversation_history},
}
],
db=db,
)
if flow_step_result["message"] != "success":
# Log the error but continue - flow step logging is not critical for user experience
log.warning(f"Failed to record agent selection flow step: {flow_step_result.get("error", "Unknown error")}")
# Continue with agent selection - user should still get their answer
step_order += 1
else:
log.error("Failed to parse the routing result. Fallback to generic agent.")
selected_agents = [AgentQuery(agent=Agents.GENERIC_AGENT, query=parsed_query)]
else:
agent_query_generic = AgentQuery(agent=Agents.GENERIC_AGENT, query=parsed_query)
selected_agents = [agent_query_generic]
flow_step_result = await upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=[
{
"step_order": step_order,
"step_type": FlowStepType.ROUTING.value,
"tool_name": None,
"content": f"Generic agent: {parsed_query}",
"execution_data": {"enhanced_conversation_history": enhanced_conversation_history},
}
],
db=db,
)
if flow_step_result["message"] != "success":
# Log the error but continue - flow step logging is not critical for user experience
log.warning(f"Failed to record agent selection result flow step: {flow_step_result.get("error", "Unknown error")}")
# Continue with returning selected agents - user should still get their answer
step_order += 1
return selected_agents, step_order, error_message
apps/pi/pi/services/chat/helpers/ask_mode_helpers.py
+739
View File
@@ -0,0 +1,739 @@
"""Helper functions for ask mode processing to keep the main flow clean."""
import asyncio
import uuid
from typing import Any
from typing import AsyncIterator
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from typing import Union
from langchain_core.messages import BaseMessage
from langchain_core.messages import HumanMessage
from langchain_core.messages import SystemMessage
from pydantic import UUID4
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import MessageMetaStepType
from pi.services.chat.utils import StandardAgentResponse
from pi.services.chat.utils import get_current_timestamp_context
from pi.services.chat.utils import mask_uuids_in_text
from pi.services.chat.utils import reasoning_dict_maker
from pi.services.chat.utils import resolve_workspace_slug
from pi.services.chat.utils import standardize_flow_step_content
from pi.services.retrievers.pg_store.message import upsert_message_flow_steps
log = logger.getChild(__name__)
# ------------------------------
# Initialization helpers
# ------------------------------
async def resolve_workspace_id_from_project(project_id: str | None) -> str | None:
"""Resolve workspace_id from project_id if workspace_id is not provided."""
if not project_id:
return None
try:
from pi.app.api.v1.helpers.plane_sql_queries import resolve_workspace_id_from_project_id
resolved_workspace_id = await resolve_workspace_id_from_project_id(project_id)
if resolved_workspace_id:
log.info(f"Resolved workspace_id {resolved_workspace_id} from project_id {project_id}")
return str(resolved_workspace_id)
except Exception as e:
log.error(f"Failed to resolve workspace_id from project_id {project_id}: {e}")
return None
def extract_or_create_query_id(user_meta: dict | None) -> uuid.UUID:
"""Extract query_id from user_meta token_id or create a new one."""
query_id = None
try:
if isinstance(user_meta, dict):
token_id = user_meta.get("token_id")
if token_id:
query_id = uuid.UUID(str(token_id))
except Exception:
pass
return query_id or uuid.uuid4()
async def resolve_workspace_slug_if_needed(workspace_id: str | None, provided_workspace_slug: str | None) -> str | None:
"""Resolve workspace slug from workspace_id if not provided."""
if provided_workspace_slug:
return provided_workspace_slug
if not workspace_id:
return None
try:
workspace_uuid = uuid.UUID(workspace_id)
return await resolve_workspace_slug(workspace_uuid, provided_workspace_slug)
except Exception as e:
log.warning(f"Failed to resolve workspace slug: {e}")
return None
async def check_and_enrich_clarification_context(
*, chat_id: UUID4, db: AsyncSession, user_meta: dict, parsed_query: str, query_id: UUID4, step_order: int
) -> Tuple[dict, int]:
"""
Check for pending clarifications and enrich user_meta with clarification context.
Returns: (enriched_user_meta, new_step_order)
"""
try:
from pi.services.retrievers.pg_store.clarifications import get_latest_pending_for_chat as _get_pending_clar
log.info(f"ChatID: {chat_id} - Checking for pending clarifications...")
clar_row = await _get_pending_clar(db=db, chat_id=uuid.UUID(str(chat_id)))
if clar_row:
clar_payload = clar_row.payload or {}
log.info(
f"ChatID: {chat_id} - Found pending clarification, enriching user_meta. Kind: {clar_row.kind}, Categories: {clar_row.categories}"
)
user_meta["clarification_context"] = {
"reason": clar_payload.get("reason"),
"missing_fields": clar_payload.get("missing_fields") or [],
"disambiguation_options": clar_payload.get("disambiguation_options") or [],
"category_hints": clar_payload.get("category_hints") or [],
"answer_text": parsed_query,
"original_query": clar_row.original_query,
"clarifies_message_id": str(clar_row.message_id),
}
log.info(f"ChatID: {chat_id} - Enriched user_meta with clarification_context")
# Record a clarification_response flow step for traceability
try:
await upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=[
{
"step_order": step_order + 1,
"step_type": FlowStepType.TOOL.value,
"tool_name": "clarification_response",
"content": standardize_flow_step_content(user_meta.get("clarification_context", {}), FlowStepType.TOOL),
"execution_data": {
"clarifies_message_id": str(clar_row.message_id),
"clarification_resolved": False,
},
}
],
db=db,
)
# Push subsequent steps by +1
step_order = step_order + 2
except Exception:
pass
else:
log.info(f"ChatID: {chat_id} - No pending clarification found")
except Exception as _e:
log.warning(f"ChatID: {chat_id} - Clarification table check failed: {_e}")
return user_meta, step_order
async def resolve_pending_clarification(*, chat_id: UUID4, db: AsyncSession, parsed_query: str, query_id: UUID4) -> None:
"""Resolve any pending clarification record for this chat."""
try:
from pi.services.retrievers.pg_store.clarifications import get_latest_pending_for_chat
from pi.services.retrievers.pg_store.clarifications import resolve_clarification
log.info(f"ChatID: {chat_id} - Processing clarification follow-up")
clar_row = await get_latest_pending_for_chat(db=db, chat_id=uuid.UUID(str(chat_id)))
if clar_row:
try:
await resolve_clarification(
db,
clarification_id=clar_row.id,
answer_text=parsed_query,
resolved_by_message_id=query_id,
)
log.info(f"ChatID: {chat_id} - Clarification resolved, continuing with tool execution")
except Exception as e:
log.warning(f"ChatID: {chat_id} - Failed to resolve clarification: {e}")
else:
log.warning(f"ChatID: {chat_id} - No clarification record found")
except Exception as e:
log.error(f"ChatID: {chat_id} - Error processing clarification follow-up: {e}")
async def handle_preset_query_flow(
chatbot_instance,
preset_query_steps: List[Dict[str, Any]],
query_id: UUID4,
chat_id: UUID4,
step_order: int,
db: AsyncSession,
user_meta: dict,
workspace_id: str,
project_id: str,
conversation_history: List[Any],
user_id: str,
query_flow_store: dict,
enhanced_query_for_processing: str,
workspace_in_context: bool,
switch_llm: str,
reasoning_container: dict,
source: Optional[str],
) -> AsyncIterator[Union[str, Dict[str, Any]]]:
"""
Handle complete preset query flow: build tools, record routing, execute, and stream response.
Returns chunks including final response marker.
"""
from pi.services.schemas.chat import ToolQuery
# Build selected tools from preset config
selected_tools = []
for step in preset_query_steps:
reasoning_messages = step.get("reasoning_messages", [])
tools_config = step.get("agents", []) # Still called "agents" in template
# Show preset reasoning messages
async for chunk in process_preset_reasoning_messages(reasoning_messages, reasoning_container):
yield chunk
# Convert preset config to ToolQuery format
for tool_config in tools_config:
tool_name = tool_config.get("agent") # Template still uses "agent" key
tool_query_str = tool_config.get("query")
if tool_name and tool_query_str:
selected_tools.append(ToolQuery(tool=tool_name, query=tool_query_str))
# Record preset routing as flow step
if selected_tools:
step_order = await record_preset_routing_step(query_id=query_id, chat_id=chat_id, step_order=step_order, selected_tools=selected_tools, db=db)
# Execute preset tool flow
if selected_tools:
async for chunk in execute_preset_tool_flow(
chatbot_instance,
preset_query_steps,
selected_tools,
user_meta,
query_id,
workspace_id,
project_id,
conversation_history,
user_id,
str(chat_id),
query_flow_store,
enhanced_query_for_processing,
step_order,
workspace_in_context,
db,
switch_llm,
reasoning_container["content"],
source,
):
yield chunk
async def execute_preset_tool_flow(
chatbot_instance,
preset_query_steps: List[Dict[str, Any]],
selected_tools: List[Any],
user_meta: dict,
query_id: UUID4,
workspace_id: str,
project_id: str,
conversation_history: List[Any],
user_id: str,
chat_id: str,
query_flow_store: dict,
enhanced_query_for_processing: str,
step_order: int,
workspace_in_context: bool,
db: AsyncSession,
switch_llm: str,
reasoning: str,
source: Optional[str],
) -> AsyncIterator[Union[str, Dict[str, Any]]]:
"""Execute preset tool flow and yield response chunks."""
try:
# For now, execute the first preset tool (templates currently define one)
preset_step = preset_query_steps[0] if len(preset_query_steps) > 0 else {}
preset_tools_cfg = preset_step.get("agents", []) if isinstance(preset_step, dict) else []
first_tool_cfg = preset_tools_cfg[0] if preset_tools_cfg else {}
# Extract preset SQL execution parameters
preset_tables = first_tool_cfg.get("tables")
preset_sql_query = first_tool_cfg.get("sql_query")
preset_placeholders = first_tool_cfg.get("placeholders_in_order")
# Align sub-step persistence for text2sql with current step
try:
if isinstance(query_flow_store, dict):
query_flow_store["step_order"] = step_order
except Exception:
pass
# Execute the preset query via the standard tool handler (handles CTEs, flow steps, entity URLs)
result = await chatbot_instance.handle_tool_query(
db=db,
tool=selected_tools[0].tool,
query=selected_tools[0].query,
user_meta=user_meta,
message_id=query_id,
workspace_id=workspace_id,
project_id=project_id,
conv_hist=conversation_history,
user_id=str(user_id),
chat_id=str(chat_id),
query_flow_store=query_flow_store,
is_multi_tool=False,
preset_tables=preset_tables,
preset_sql_query=preset_sql_query,
preset_placeholders=preset_placeholders,
)
# Persist a high-level tool execution step so downstream views can extract entity_urls and summaries
try:
tool_name_value = getattr(selected_tools[0].tool, "value", str(selected_tools[0].tool))
flow_step_payload = {
"step_order": step_order,
"step_type": FlowStepType.TOOL.value,
"tool_name": tool_name_value,
"content": standardize_flow_step_content(result, FlowStepType.TOOL),
"execution_data": result if isinstance(result, dict) else {"results": str(result)},
"is_planned": False,
"is_executed": True,
"execution_success": "success",
}
await upsert_message_flow_steps(message_id=query_id, chat_id=uuid.UUID(str(chat_id)), flow_steps=[flow_step_payload], db=db)
step_order += 1
except Exception as e:
log.warning(f"ChatID: {chat_id} - Failed to persist preset tool execution step: {e}")
# Prepare streaming of the final user-visible text via LLM combiner (applies phrasing rules)
results_text = StandardAgentResponse.extract_results(result) # plain tool output
try:
# Include attachments if any are in shared context
attachment_blocks = getattr(chatbot_instance, "get_current_attachment_blocks", lambda: None)()
except Exception:
attachment_blocks = None
base_stream = chatbot_instance.combined_response_stream(
query=(enhanced_query_for_processing or selected_tools[0].query),
responses=results_text,
conversation_history=conversation_history,
user_meta=user_meta or {},
user_id=str(user_id),
attachment_blocks=attachment_blocks,
workspace_in_context=workspace_in_context,
)
# Stream and persist assistant message
final_response = ""
async for chunk in chatbot_instance._process_response(
base_stream, chat_id, query_id, uuid.uuid4(), switch_llm, db, reasoning=reasoning, source=source
):
if chunk.startswith("__FINAL_RESPONSE__"):
final_response = chunk[len("__FINAL_RESPONSE__") :]
continue
yield chunk
if final_response:
query_flow_store["answer"] = final_response
except Exception as e:
log.error(f"ChatID: {chat_id} - Error executing preset flow: {e}")
raise
async def record_preset_routing_step(*, query_id: UUID4, chat_id: UUID4, step_order: int, selected_tools: List[Any], db: AsyncSession) -> int:
"""Record preset routing as flow step with timeout protection."""
routing_content = [{"tool": tool_query.tool.name, "query": tool_query.query} for tool_query in selected_tools]
try:
flow_step_result = await asyncio.wait_for(
upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=[
{
"step_order": step_order,
"step_type": FlowStepType.ROUTING.value,
"tool_name": None,
"content": standardize_flow_step_content(routing_content, FlowStepType.ROUTING),
"execution_data": {},
}
],
db=db,
),
timeout=2.0,
)
if flow_step_result["message"] != "success":
log.warning(f"Failed to record RAG routing flow step: {flow_step_result.get("error", "Unknown error")}")
except asyncio.TimeoutError:
log.warning("Timed out recording RAG routing flow step; continuing")
except Exception as e:
log.warning(f"Failed to record RAG routing flow step: {e}")
return step_order + 1
async def process_preset_reasoning_messages(reasoning_messages: List[str], reasoning_container: dict) -> AsyncIterator[Union[str, Dict[str, Any]]]:
"""Yield preset reasoning messages as special reasoning blocks."""
for reasoning_msg in reasoning_messages:
masked_reasoning_msg = mask_uuids_in_text(reasoning_msg)
stage = "ask_preset_reasoning"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content=masked_reasoning_msg)
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
async def collect_conversation_attachments(*, chat_id: UUID4, db: AsyncSession, current_attachment_ids: List[UUID4]) -> List[UUID4]:
"""Collect all attachments from conversation history and merge with current attachments."""
all_attachment_ids = current_attachment_ids.copy() if current_attachment_ids else []
try:
from pi.services.retrievers.pg_store.chat import retrieve_chat_history
raw_history = await retrieve_chat_history(chat_id=chat_id, db=db, dialogue_object=True)
if raw_history.get("dialogue"):
attachment_count = 0
for qa_pair in raw_history["dialogue"]:
if qa_pair.get("attachments"):
for att in qa_pair["attachments"]:
att_id = att.get("id")
if att_id and att_id not in [str(aid) for aid in all_attachment_ids]:
all_attachment_ids.append(uuid.UUID(att_id))
attachment_count += 1
if attachment_count > 0:
log.info(f"ChatID: {chat_id} - Collected {attachment_count} attachments from conversation history")
except Exception as e:
log.warning(f"ChatID: {chat_id} - Failed to collect conversation attachments: {e}")
return all_attachment_ids
async def generate_chat_title_if_needed(
*,
chatbot_instance,
chat_id: UUID4,
query_id: UUID4,
db: AsyncSession,
final_response: str,
parsed_query: str | None,
query: str,
) -> None:
"""
Generate chat title for new chats (first message only).
Checks if this is the first message and generates a title either:
- Directly from the query (for error responses)
- Using LLM (for successful responses)
"""
try:
from sqlalchemy import select
from pi.app.models import Message
# Get the message sequence number to determine if this is a new chat
stmt = select(Message).where(Message.id == query_id) # type: ignore[arg-type]
result = await db.execute(stmt)
message = result.scalar_one_or_none()
if message:
message_sequence = message.sequence
is_first_message = message_sequence == 1
else:
log.warning(f"ChatID: {chat_id} - Message {query_id} not found for sequence check")
is_first_message = False
if is_first_message:
try:
# For error scenarios, use the user's question directly as title
is_error_response = (
not final_response
or final_response.startswith("An unexpected error occurred")
or final_response == "TOOL_ORCHESTRATION_FAILURE"
or final_response.startswith("Action execution is not available")
or final_response.startswith("Chat ID is required")
or final_response.startswith("Your request timed out")
)
title = None
if is_error_response:
# Use the user's original question as title, with some cleanup
title = (parsed_query if parsed_query is not None else query).strip()
# Truncate if too long (keeping it under ~60 characters for readability)
if len(title) > 60:
title = title[:57] + "..."
await chatbot_instance.set_chat_title_directly(chat_id, title, db)
elif final_response:
# For successful responses, generate title using LLM
chat_history = [parsed_query if parsed_query is not None else query, final_response]
# Ensure both elements are strings
chat_history = [str(item) for item in chat_history if item is not None]
if len(chat_history) >= 2: # Make sure we have both query and response
title = await chatbot_instance.generate_title(chat_id, chat_history, db)
# Generated chat title search index upserted via Celery background task
else:
log.warning(f"ChatID: {chat_id} - Not enough valid content to generate title")
title = (parsed_query if parsed_query is not None else query).strip()
# Truncate if too long (keeping it under ~60 characters for readability)
if len(title) > 60:
title = title[:57] + "..."
await chatbot_instance.set_chat_title_directly(chat_id, title, db)
else:
log.warning(f"ChatID: {chat_id} - No final_response available for title generation")
if title:
log.info(f"ChatID: {chat_id} - Generated title: {title}")
except Exception as e:
log.error(f"Error generating title: {str(e)}")
except Exception as e:
log.error(f"Error checking message sequence for title generation: {str(e)}")
def log_ask_mode_request_details(data: Any, context: Dict[str, Any]) -> None:
"""
Log all important request details for ask mode processing.
Args:
data: ChatRequest object with original request fields
context: Dict with derived/computed values:
- enhanced_query_for_processing
- attachment_context
- workspace_slug (resolved)
- workspace_id (resolved)
- workspace_in_context (computed)
"""
chat_id = data.chat_id
log.info(f"ChatID: {chat_id} - Input query: {data.query}")
log.info(f"ChatID: {chat_id} - Enhanced query: {context.get("enhanced_query_for_processing")}")
log.info(f"ChatID: {chat_id} - Attachment context: {context.get("attachment_context")}")
log.info(f"ChatID: {chat_id} - User meta: {data.context}")
log.info(f"ChatID: {chat_id} - Workspace in context: {context.get("workspace_in_context", data.workspace_in_context)}")
log.info(f"ChatID: {chat_id} - Workspace slug: {context.get("workspace_slug", data.workspace_slug)}")
log.info(f"ChatID: {chat_id} - Workspace ID: {context.get("workspace_id", str(data.workspace_id) if data.workspace_id else None)}")
log.info(f"ChatID: {chat_id} - Is New Chat: {data.is_new}")
log.info(f"ChatID: {chat_id} - Source: {data.source}")
log.info(f"ChatID: {chat_id} - Is Project Chat: {data.is_project_chat}")
log.info(f"ChatID: {chat_id} - User ID: {data.user_id}")
log.info(f"ChatID: {chat_id} - Project ID: {str(data.project_id) if data.project_id else None}")
log.info(f"ChatID: {chat_id} - LLM: {data.llm}")
async def construct_enhanced_prompt_and_context(
enhanced_conversation_history: str | None,
project_id: str | None,
workspace_id: str | None,
user_id: str | None,
user_meta: dict | None,
workspace_in_context: bool = True,
) -> Tuple[str, str]:
"""
Build enhanced custom prompt and context block for LLM orchestration.
Returns:
Tuple of (custom_prompt, context_block)
"""
from pi.services.chat.helpers.tool_utils import build_clarification_context_block
from pi.services.chat.utils import get_current_timestamp_context
custom_prompt = ""
# Inject enhanced conversation history if available to guide orchestration
try:
if enhanced_conversation_history and isinstance(enhanced_conversation_history, str) and enhanced_conversation_history.strip():
custom_prompt += f"\n\n**CONVERSATION HISTORY & ACTION CONTEXT:**\n{enhanced_conversation_history}\n"
except Exception:
pass
# Provide PROJECT/USER/TIME context to the tool planner
context_block = ""
try:
# Check if workspace is in context
if not workspace_in_context:
# No workspace context - guide LLM on how to handle queries
context_block += "\n\n**⚠️ NO WORKSPACE CONTEXT:**\n\n**Instructions:**\n- For general questions about Plane features and functionality, use docs_search_tool to search documentation\n- For questions NOT related to Plane, answer directly using your knowledge\n- For questions about specific workspace data (issues, projects, cycles, team members, etc.), politely inform the user:\n 'To access your workspace data, please enable workspace context using the chat input box options.'\n" # noqa: E501
elif project_id:
# Project scoping (project chat)
context_block += f"\n\n**🔥 PROJECT CONTEXT (CRITICAL):**\nProject ID: {project_id} ⚠️ (FOR YOUR USE ONLY - DO NOT SHOW THIS UUID TO USER)\n\n**IMPORTANT SCOPING RULES:**\n- This is a PROJECT-LEVEL chat - ALL operations are scoped to THIS PROJECT ONLY\n- When the request mentions 'current cycle', 'current module', 'work items', etc. - it means ONLY within THIS PROJECT\n- Use this project_id for ALL tools that accept project_id parameter\n- DO NOT query across all projects - scope everything to THIS specific project\n- User refers to 'this project'/'the project'/'current project' = use this project_id" # noqa: E501
else:
# Workspace-level context (no specific project)
context_block += f"\n\n**🌐 WORKSPACE CONTEXT (CRITICAL):**\nWorkspace ID: {workspace_id} ⚠️ (FOR YOUR USE ONLY - DO NOT SHOW THIS UUID TO USER)\n\n**IMPORTANT SCOPING RULES:**\n- This is a WORKSPACE-LEVEL chat - queries can span MULTIPLE PROJECTS\n- When the request mentions 'last cycle', 'this cycle', 'work items', etc. WITHOUT specifying a project - it could be in ANY project\n- Use list_member_projects (without limit or with high limit) to get ALL projects in the workspace\n- Then iterate through projects to find relevant cycles/modules/work-items\n- CRITICAL: Do NOT limit to 1 project unless the user specifically names or refers to a specific project" # noqa: E501
# User context
if user_id:
context_block += (
"\n**USER CONTEXT:**\n"
f"User ID: {user_id} ⚠️ (FOR YOUR USE ONLY - DO NOT SHOW THIS UUID TO USER)\n"
"- Use this to resolve pronouns like 'me', 'my', 'assigned to me'.\n"
)
# Date/time context
dt_ctx = None
try:
if user_id:
dt_ctx = await get_current_timestamp_context(user_id)
except Exception:
dt_ctx = None
if dt_ctx:
custom_prompt += f"\n{dt_ctx}\n"
except Exception:
# Non-fatal; continue without context block if any error occurs
pass
# CRITICAL: Inject clarification context if this is a follow-up to ask_for_clarification
if user_meta and isinstance(user_meta, dict) and user_meta.get("clarification_context"):
custom_prompt += build_clarification_context_block(user_meta.get("clarification_context"))
return custom_prompt, context_block
async def resolve_workspace_id_and_slug(
*,
workspace_id: str | None,
project_id: str | None,
chat_id: UUID4,
) -> Tuple[str | None, str | None]:
"""
Resolve workspace_id from project_id if needed, then get workspace_slug.
Returns:
Tuple of (resolved_workspace_id, workspace_slug)
"""
from pi.app.api.v1.helpers.plane_sql_queries import get_workspace_slug
from pi.app.api.v1.helpers.plane_sql_queries import resolve_workspace_id_from_project_id
# Resolve workspace_id from project_id if needed (for project-level chats)
if not workspace_id and project_id:
try:
resolved_workspace_id = await resolve_workspace_id_from_project_id(str(project_id))
if resolved_workspace_id:
workspace_id = str(resolved_workspace_id)
log.info(f"ChatID: {chat_id} - Resolved workspace_id {workspace_id} from project_id {project_id} (Phase 1)")
except Exception as e:
log.error(f"ChatID: {chat_id} - Failed to resolve workspace_id from project_id {project_id}: {e}")
# Build workspace slug
workspace_slug = None
if workspace_id:
try:
workspace_slug = await get_workspace_slug(workspace_id)
except Exception as e:
log.warning(f"ChatID: {chat_id} - Failed to get workspace_slug for workspace_id {workspace_id}: {e}")
return workspace_id, workspace_slug
# ------------------------------
# Fast-path helpers (no workspace context)
# ------------------------------
async def create_fast_path_stream(
chatbot_instance,
query_id: UUID4,
db: AsyncSession,
chat_id: UUID4,
user_meta: Dict[str, Any] | None,
user_id: str | None,
enhanced_query_for_processing: str,
enhanced_conversation_history: str | None,
reasoning_container: Dict[str, str] | None = None,
) -> AsyncIterator[Union[str, Dict[str, Any]]]:
"""
Create a fast-path stream for queries without workspace context.
This bypasses tool orchestration and directly calls the LLM with a simple prompt.
Used when workspace_in_context is False for both Ask and Build modes.
Args:
chatbot_instance: The chatbot instance with tool_llm attribute
query_id: Message ID for token tracking
db: Database session
chat_id: Chat ID for logging
user_meta: User metadata dict (may contain first_name, email)
user_id: User ID for timestamp context
enhanced_query_for_processing: The user's query (may include attachment context)
enhanced_conversation_history: Previous conversation history
Yields:
str: Streaming chunks from LLM response
"""
log.info(f"ChatID: {chat_id} - No workspace context, using fast path (direct LLM call)")
# Extract user name from user_meta if available
user_first_name = None
user_email = None
if user_meta and isinstance(user_meta, dict):
user_first_name = user_meta.get("first_name")
user_email = user_meta.get("email")
# Get timestamp context
dt_ctx = None
try:
if user_id:
dt_ctx = await get_current_timestamp_context(user_id)
except Exception:
pass
# Build context block
context_parts = []
if user_first_name:
context_parts.append(f"**User:** {user_first_name}")
elif user_email:
context_parts.append(f"**User:** {user_email}")
if dt_ctx:
context_parts.append(dt_ctx.strip())
context_block = "\n".join(context_parts) if context_parts else ""
# Build system prompt
system_prompt = f"""You are Plane AI, an AI assistant for Plane - a work management platform.
**Current Mode:** No workspace context available.
{context_block}
**Your role:**
- Answer general knowledge questions not related to Plane
- If user asks about their specific workspace data (issues, projects, cycles, team members, etc.), politely inform them:
"To access your workspace data, please enable workspace context using the chat input box options."
Be helpful, concise, and professional."""
# Build message list
messages: List[BaseMessage] = [SystemMessage(content=system_prompt)]
if enhanced_conversation_history and isinstance(enhanced_conversation_history, str) and enhanced_conversation_history.strip():
messages.append(SystemMessage(content=f"Conversation History:\n{enhanced_conversation_history}"))
messages.append(HumanMessage(content=enhanced_query_for_processing))
# Setup LLM with token tracking
llm = chatbot_instance.tool_llm
llm.set_tracking_context(query_id, db, MessageMetaStepType.TOOL_ORCHESTRATION, chat_id=str(chat_id))
# Signal to the caller that we're generating the final response
# This marker is used by chat.py to start collecting chunks for DB persistence
stage = "final_response"
reasoning_chunk_dict = reasoning_dict_maker(stage=stage, tool_name="", tool_query="", content="")
if reasoning_container is not None:
reasoning_container["content"] += reasoning_chunk_dict["header"] + reasoning_chunk_dict["content"]
yield reasoning_chunk_dict
# Stream LLM response
# Marker used by collector to start capturing final response
async for chunk in llm.astream(messages):
if getattr(chunk, "content", None):
yield str(chunk.content)
apps/pi/pi/services/chat/helpers/batch_action_orchestrator.py
-764
View File
@@ -1,764 +0,0 @@
"""LLM-based orchestration for batch action execution."""
import re
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from langchain_core.messages import SystemMessage
from langchain_core.messages import ToolMessage
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.services.actions.tools.entity_search import get_entity_search_tools
from pi.services.chat.chat import PlaneChatBot
from .batch_execution_context import BatchExecutionContext
from .batch_execution_helpers import build_execution_prompt
from .batch_execution_helpers import extract_entity_info
from .batch_execution_helpers import update_flow_step_execution_status
log = logger.getChild(__name__)
class BatchActionOrchestrator:
"""Orchestrates batch action execution using LLM."""
def __init__(self, chatbot: PlaneChatBot, db: AsyncSession):
self.chatbot = chatbot
self.db = db
async def execute_planned_actions(
self,
context: BatchExecutionContext,
original_query: str,
planned_actions: List[Dict[str, Any]],
access_token: str,
workspace_slug: str,
) -> BatchExecutionContext:
"""Execute all planned actions using LLM orchestration."""
context.set_total_planned(len(planned_actions))
try:
# Build execution prompt
execution_prompt = build_execution_prompt(original_query, planned_actions)
# Create method executor and tools
_, combined_tools = await self._setup_execution_tools(access_token, workspace_slug, context, planned_actions) # noqa: E501
if not combined_tools:
context.add_execution_failure("setup", "tool_setup", "No execution tools available", None)
return context
# Execute using LLM orchestration
await self._orchestrate_with_llm(execution_prompt, combined_tools, planned_actions, context)
# Mark execution as completed
context.mark_completed()
except Exception as e:
log.error(f"Batch execution failed: {e}")
context.add_execution_failure("orchestration", "system_error", str(e), None)
return context
async def _setup_execution_tools(
self, access_token: str, workspace_slug: str, context: BatchExecutionContext, planned_actions: List[Dict[str, Any]]
):
"""Set up the execution tools for the planned actions."""
from pi import settings
from pi.services.actions import MethodExecutor
from pi.services.actions import PlaneActionsExecutor
# Create actions executor
if access_token.startswith("plane_api_"):
actions_executor = PlaneActionsExecutor(api_key=access_token, base_url=settings.plane_api.HOST)
else:
actions_executor = PlaneActionsExecutor(access_token=access_token, base_url=settings.plane_api.HOST)
method_executor = MethodExecutor(actions_executor)
# Build context for tools - extract project_id from planned actions
project_id = None
for action in planned_actions:
# Look for project_id in action args
args = action.get("args", {})
if args.get("project_id"):
project_id = args["project_id"]
break
if not project_id:
log.warning("🔧 WARNING: No project_id found in planned actions - context auto-fill may fail")
else:
# Check if project_id needs resolution (is not a UUID)
import re
uuid_regex = re.compile(r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")
if not uuid_regex.match(project_id):
# log.info(f"🔧 Resolving project_id '{project_id}' to UUID before building context")
# Create temporary tools for resolution
temp_executor = PlaneActionsExecutor(access_token=access_token, base_url=settings.plane_api.HOST)
temp_method_executor = MethodExecutor(temp_executor)
temp_context = {"workspace_slug": workspace_slug}
temp_tools = self.chatbot._build_method_tools("entity_search", temp_method_executor, temp_context)
# Resolve project_id
resolved_project_id = await self._resolve_entity_id("project", project_id, temp_tools, None)
if resolved_project_id:
# log.info(f"🔧 Successfully resolved project_id '{project_id}' to '{resolved_project_id}'")
project_id = resolved_project_id
else:
log.error(f"🔧 Failed to resolve project_id '{project_id}' - will use as-is")
# workspace_slug is now auto-filled in tool implementations (hidden from LLM)
tool_context: Dict[str, Any] = {
"workspace_id": str(context.workspace_id),
"workspace_slug": workspace_slug,
"project_id": project_id, # 🔥 FIX: Extract from planned actions
"user_id": str(context.user_id),
"conversation_history": [],
"user_meta": {},
"is_project_chat": context.is_project_chat, # 🔥 FIX: Pass is_project_chat to tools
}
log.info(
f"🔧 EXECUTION TOOL CONTEXT: is_project_chat={context.is_project_chat}, workspace_id={context.workspace_id}, chat_id={context.chat_id}"
)
# Determine which categories of tools we need based on planned actions
categories = set()
for action in planned_actions:
tool_name = action.get("tool_name", "")
# Ensure tool_name is always a string to avoid NoneType iteration errors
if tool_name is None:
tool_name = ""
elif not isinstance(tool_name, str):
tool_name = str(tool_name)
if tool_name and "_" in tool_name:
category = tool_name.split("_")[0]
# Map epics to workitems category since epics tools are part of workitems
if category == "epics":
category = "workitems"
categories.add(category)
log.info(f"🔧 Building execution tools for categories: {categories}")
# Build tools for all required categories
combined_tools = []
tool_names_seen = set()
for category in categories:
try:
tools_for_category = self.chatbot._build_method_tools(category, method_executor, tool_context)
# Only add tools we haven't seen before (deduplicate by name)
tool_names_list = []
for tool in tools_for_category:
if tool.name not in tool_names_seen:
combined_tools.append(tool)
tool_names_seen.add(tool.name)
tool_names_list.append(tool.name)
log.info(f"🔧 Built tools for category '{category}': {tool_names_list}")
except Exception as e:
log.warning(f"Failed to build tools for category {category}: {e}")
# Add entity search tools for placeholder resolution (only if not already included)
try:
entity_search_tools = get_entity_search_tools(method_executor, tool_context)
for tool in entity_search_tools:
if tool.name not in tool_names_seen:
combined_tools.append(tool)
tool_names_seen.add(tool.name)
log.info(
f"🔧 Added {len([t for t in entity_search_tools if t.name not in tool_names_seen])} entity search tools for placeholder resolution"
)
except Exception as e:
log.warning(f"Failed to build entity search tools: {e}")
return method_executor, combined_tools
async def _orchestrate_with_llm(
self,
execution_prompt: str,
combined_tools: List,
planned_actions: List[Dict[str, Any]],
context: BatchExecutionContext,
):
"""Use LLM to orchestrate the execution of planned actions."""
# Initialize conversation
messages: List[Any] = [SystemMessage(content=execution_prompt)]
# Bind tools to LLM
llm_with_tools = self.chatbot.tool_llm.bind_tools(combined_tools)
# Set tracking context for action execution LLM calls
from pi.app.models.enums import MessageMetaStepType
llm_with_tools.set_tracking_context(context.message_id, self.db, MessageMetaStepType.ACTION_EXECUTION, chat_id=str(context.chat_id))
# Track which actions have been executed
executed_tool_names = set()
used_step_ids: set = set() # Track which step_ids have been matched to prevent double-matching
max_iterations = 15 # Safety limit
iteration_count = 0
# Start the orchestration conversation
response = await llm_with_tools.ainvoke(messages)
# Continue until no more tool calls or we hit limits
while (
hasattr(response, "tool_calls") and getattr(response, "tool_calls", None) and iteration_count < max_iterations and context.can_continue()
):
# Add the response with tool calls to messages first
messages.append(response)
iteration_count += 1
tool_messages = []
execution_failed = False
tool_calls = getattr(response, "tool_calls", [])
# Ensure tool_calls is always a list
if tool_calls is None:
tool_calls = []
elif not isinstance(tool_calls, (list, tuple)):
log.warning(f"Unexpected tool_calls type: {type(tool_calls)}, converting to list")
tool_calls = [tool_calls] if tool_calls else []
# Execute all tool calls in this round
for tool_call in tool_calls:
try:
# Extract tool call information - handle both LangChain objects and dictionaries
if isinstance(tool_call, dict):
# Handle dictionary format (what we're actually getting)
tool_name = tool_call.get("name", "")
tool_args = tool_call.get("args", {})
tool_id = tool_call.get("id", "")
else:
# Handle LangChain object format
tool_name = getattr(tool_call, "name", "")
tool_args = getattr(tool_call, "args", {})
tool_id = getattr(tool_call, "id", "")
if not tool_name:
continue
if not tool_id:
continue
# Ensure tool_name is a string and not None
tool_name = str(tool_name) if tool_name is not None else ""
if not tool_name:
continue
# Find the corresponding planned action with argument matching and consumption tracking
planned_action = self._find_planned_action(tool_name, planned_actions, tool_args, used_step_ids)
if not planned_action:
# Try to find by partial match (LLM might use slightly different names)
planned_action = self._find_planned_action_partial(tool_name, planned_actions)
# Debug logging for step_id matching
if not planned_action:
log.warning(f"⚠️ Could not match tool_call {tool_name} to any planned action")
# Ensure we have a valid step_id
if planned_action and planned_action.get("step_id"):
step_id = planned_action["step_id"]
else:
step_id = f"unknown_{context.current_step}"
log.warning(f"No planned action found for tool {tool_name}, using fallback step_id: {step_id}")
# DEFENSIVE FIX: Before placeholder resolution, check if action_summary has project.id
# and tool_args.project_id is missing or non-UUID. Extract from summary if available.
uuid_pattern = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.IGNORECASE)
if planned_action:
action_summary = planned_action.get("action_summary", {})
log.info(
f"DEBUG: action_summary type={type(action_summary)}, keys={action_summary.keys() if isinstance(action_summary, dict) else "N/A"}" # noqa: E501
)
if isinstance(action_summary, dict):
parameters = action_summary.get("parameters", {})
log.info(
f"DEBUG: parameters type={type(parameters)}, keys={parameters.keys() if isinstance(parameters, dict) else "N/A"}" # noqa: E501
)
if isinstance(parameters, dict):
project_block = parameters.get("project", {})
log.info(f"DEBUG: project_block type={type(project_block)}, content={project_block}") # noqa: E501
if isinstance(project_block, dict):
summary_project_id = project_block.get("id")
log.info(
f"DEBUG: summary_project_id={summary_project_id}, is_uuid={uuid_pattern.match(summary_project_id) if isinstance(summary_project_id, str) else False}" # noqa: E501
)
if isinstance(summary_project_id, str) and uuid_pattern.match(summary_project_id):
# Check if tool_args needs this project_id
current_project_id = tool_args.get("project_id") if isinstance(tool_args, dict) else None
log.info(
f"DEBUG: current_project_id={current_project_id}, is_uuid={uuid_pattern.match(current_project_id) if isinstance(current_project_id, str) else False}" # noqa: E501
)
if not current_project_id or (
isinstance(current_project_id, str) and not uuid_pattern.match(current_project_id)
):
# Set project_id from action_summary to avoid resolution
if isinstance(tool_args, dict):
tool_args["project_id"] = summary_project_id
log.info(f"✅ Extracted project_id from action_summary for {tool_name}: {summary_project_id}")
# Check if tool_args contain placeholders that need resolution
if self._has_placeholders(tool_args):
# log.info(f"Detected placeholders in {tool_name} arguments: {tool_args}")
# Force placeholder resolution before execution
resolved_args = await self._resolve_placeholders(tool_args, combined_tools, context)
if resolved_args:
# log.info(f"Successfully resolved placeholders for {tool_name}: {resolved_args}")
tool_args = resolved_args
else:
# Placeholder resolution failed
error_msg = f"Failed to resolve placeholders in {tool_name} arguments: {tool_args}"
artifact_id = planned_action.get("artifact_id") if planned_action else None
planned_sequence = planned_action.get("step_order") if planned_action else None
artifact_type = planned_action.get("artifact_type") if planned_action else None
context.add_execution_failure(
step_id, tool_name, error_msg, artifact_id, sequence=planned_sequence, artifact_type=artifact_type
)
execution_failed = True
continue
else:
log.debug(f"No placeholders detected in {tool_name} arguments: {tool_args}")
# Execute the tool
result = await self._execute_tool(tool_name, tool_args, combined_tools)
# Check if execution was successful
is_success = self._check_execution_success(result)
if is_success:
# Extract entity information
entity_info = extract_entity_info(str(result), tool_name)
# Coerce to a single entity dict for downstream consumers (type-safe for mypy)
single_entity_info: Optional[Dict[str, Any]]
if isinstance(entity_info, list):
single_entity_info = entity_info[0] if entity_info else None
elif isinstance(entity_info, dict):
single_entity_info = entity_info
else:
single_entity_info = None
# Record successful execution
artifact_id = planned_action.get("artifact_id") if planned_action else None
planned_sequence = planned_action.get("step_order") if planned_action else None
artifact_type = planned_action.get("artifact_type") if planned_action else None
context.add_execution_result(
step_id, tool_name, result, single_entity_info, artifact_id, sequence=planned_sequence, artifact_type=artifact_type
)
executed_tool_names.add(tool_name)
# Update flow step in database
if planned_action and planned_action.get("step_id"):
# For manual edits, pass the version_id through
version_id = planned_action.get("version_id")
await update_flow_step_execution_status(step_id, str(result), entity_info, True, self.db, version_id)
else:
# Record failed execution
error_msg = f"Tool execution failed: {result}"
artifact_id = planned_action.get("artifact_id") if planned_action else None
planned_sequence = planned_action.get("step_order") if planned_action else None
artifact_type = planned_action.get("artifact_type") if planned_action else None
context.add_execution_failure(
step_id, tool_name, error_msg, artifact_id, sequence=planned_sequence, artifact_type=artifact_type
)
# Update flow step in database
if planned_action and planned_action.get("step_id"):
# For manual edits, pass the version_id through (for failed cases too)
version_id = planned_action.get("version_id")
await update_flow_step_execution_status(step_id, str(result), None, False, self.db, version_id)
# Mark execution as failed but continue to create tool message
execution_failed = True
# Create tool message for conversation continuity (regardless of success/failure)
tool_message = ToolMessage(content=str(result), tool_call_id=tool_id)
tool_messages.append(tool_message)
except Exception as e:
log.error(f"Error executing tool {tool_name}: {e}")
# Use fallback step_id for error cases
error_step_id = step_id if "step_id" in locals() and step_id is not None else f"unknown_{context.current_step}"
artifact_id = planned_action.get("artifact_id") if planned_action else None
planned_sequence = planned_action.get("step_order") if planned_action else None
artifact_type = planned_action.get("artifact_type") if planned_action else None
context.add_execution_failure(
error_step_id, tool_name, str(e), artifact_id, sequence=planned_sequence, artifact_type=artifact_type
)
# Create error tool message
tool_message = ToolMessage(content=f"Error executing {tool_name}: {str(e)}", tool_call_id=tool_id)
tool_messages.append(tool_message)
execution_failed = True
# Add tool results to conversation (this must happen before any break)
messages.extend(tool_messages)
# Verify that all tool calls in this response have corresponding tool messages
response_tool_call_ids = set()
for tc in tool_calls:
if isinstance(tc, dict):
tool_id = tc.get("id", "")
else:
tool_id = getattr(tc, "id", "")
if tool_id:
response_tool_call_ids.add(tool_id)
added_tool_message_ids = {getattr(tm, "tool_call_id", "") for tm in tool_messages if hasattr(tm, "tool_call_id")}
missing_in_this_round = response_tool_call_ids - added_tool_message_ids
if missing_in_this_round:
# Create missing tool messages for this round
for missing_id in missing_in_this_round:
placeholder_message = ToolMessage(content="Tool execution was interrupted or failed", tool_call_id=missing_id)
messages.append(placeholder_message)
# Check if we should continue
if not context.can_continue():
break
# Check if all planned actions have been executed (by step_id count, not tool name uniqueness)
total_planned_actions = len(planned_actions)
total_executed_actions = len(used_step_ids)
if total_executed_actions >= total_planned_actions:
# log.info(f"All {total_planned_actions} planned actions have been attempted. Stopping execution loop.")
break
# If execution failed, stop here but tool messages are already added
if execution_failed:
break
# Get next response from LLM
try:
response = await llm_with_tools.ainvoke(messages)
except Exception as e:
log.error(f"Error getting LLM response: {e}")
context.add_execution_failure("llm", "orchestration_error", str(e), None)
break
def _find_planned_action(
self, tool_name: str, planned_actions: List[Dict[str, Any]], tool_args: Optional[Dict[str, Any]] = None, used_step_ids: Optional[set] = None
) -> Optional[Dict[str, Any]]:
"""Find the planned action that matches the tool name and hasn't been used yet."""
if used_step_ids is None:
used_step_ids = set()
# First try to match by tool name and key arguments for better accuracy
if tool_args:
for action in planned_actions:
artifact_id = action.get("artifact_id")
if action.get("tool_name") == tool_name and artifact_id not in used_step_ids:
# Check if key arguments match (for better matching when there are multiple identical tool calls)
planned_args = action.get("args", {})
# For workitem operations, match on issue_id
if tool_name.startswith("workitems_") and "issue_id" in tool_args and "issue_id" in planned_args:
if tool_args["issue_id"] == planned_args["issue_id"]:
if artifact_id: # Use artifact_id for tracking instead of step_id
used_step_ids.add(artifact_id)
return action
# For create operations with identical tool names, match on core identifying fields
elif tool_name.endswith("_create") and self._create_args_match(tool_args, planned_args):
if artifact_id: # Use artifact_id for tracking instead of step_id
used_step_ids.add(artifact_id)
return action
# For other operations, can add more specific matching logic as needed
elif self._args_match(tool_args, planned_args):
if artifact_id: # Use artifact_id for tracking instead of step_id
used_step_ids.add(artifact_id)
return action
# Fallback: match by tool name only, but skip already used actions
for action in planned_actions:
artifact_id = action.get("artifact_id")
if action.get("tool_name") == tool_name and artifact_id not in used_step_ids:
if artifact_id: # Use artifact_id for tracking
used_step_ids.add(artifact_id)
return action
return None
def _create_args_match(self, tool_args: Dict[str, Any], planned_args: Dict[str, Any]) -> bool:
"""Check if tool arguments match for create operations (where unique identifying fields like name are key)."""
if not tool_args and not planned_args:
return True
# For create operations, the 'name' field is usually the most distinctive identifier
# If both have name fields, they must match exactly for create operations
if "name" in tool_args and "name" in planned_args:
if str(tool_args["name"]) != str(planned_args["name"]):
return False
# If names match, check other fields for additional validation
other_fields = ["project_id", "priority", "state_id", "type_id"]
for field in other_fields:
if field in tool_args and field in planned_args:
if str(tool_args[field]) != str(planned_args[field]):
return False
return True
# If no name field, fall back to matching on other distinctive fields
key_fields = ["project_id", "priority", "state_id", "type_id"]
matches = 0
total_fields = 0
for field in key_fields:
if field in tool_args and field in planned_args:
total_fields += 1
if str(tool_args[field]) == str(planned_args[field]):
matches += 1
# Require high match ratio for create operations
return total_fields > 0 and (matches / total_fields) >= 0.8
def _args_match(self, tool_args: Dict[str, Any], planned_args: Dict[str, Any]) -> bool:
"""Check if tool arguments substantially match planned arguments."""
if not tool_args and not planned_args:
return True
# Check key fields that should match
key_fields = ["issue_id", "project_id", "module_id", "cycle_id", "label_id", "state_id", "user_id"]
for field in key_fields:
if field in tool_args and field in planned_args:
if str(tool_args[field]) != str(planned_args[field]):
return False
return True
def _find_planned_action_partial(self, tool_name: str, planned_actions: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
"""Find a planned action that matches the tool name partially."""
for action in planned_actions:
planned_tool_name = action.get("tool_name", "")
# Ensure both tool names are strings before string operations
if planned_tool_name and isinstance(planned_tool_name, str) and tool_name and isinstance(tool_name, str):
if tool_name in planned_tool_name or planned_tool_name in tool_name:
return action
return None
async def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any], tools: List) -> Any:
"""Execute a specific tool with the given arguments."""
# Ensure tool_name is a string
if not tool_name or not isinstance(tool_name, str):
log.error(f"Invalid tool name: {tool_name}")
return f"Error: Invalid tool name: {tool_name}"
# Find the tool
tool_func = next((t for t in tools if getattr(t, "name", "") == tool_name), None)
if not tool_func:
log.error(f"Tool {tool_name} not found")
return f"Error: Tool {tool_name} not found"
try:
# Execute the tool
if hasattr(tool_func, "ainvoke"):
result = await tool_func.ainvoke(tool_args)
else:
result = tool_func.invoke(tool_args)
return result
except Exception as e:
log.error(f"Error executing tool {tool_name}: {e}")
return f"Error executing {tool_name}: {str(e)}"
def _check_execution_success(self, result: Any) -> bool:
"""Check if a tool execution was successful based on the result."""
result_str = str(result).lower()
# First check for explicit success indicators (higher priority)
success_indicators = ["✅", "successfully", "created", "updated", "added"]
for indicator in success_indicators:
if indicator in result_str:
return True
# Then check for failure indicators (but be more specific to avoid false positives)
failure_indicators = ["❌", "failed", "error:", "bad request", "not found", "invalid"]
for indicator in failure_indicators:
if indicator in result_str:
return False
# Default to success if no clear indicators
return True
def _has_placeholders(self, tool_args: Dict[str, Any]) -> bool:
"""Check if tool arguments contain placeholders or non-UUID *_id values that need resolution."""
import re
log.debug(f"_has_placeholders input args: {tool_args}")
# Regex for matching standard UUIDs - fixed pattern to handle UUIDs properly
uuid_regex = re.compile(r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")
entity_keys = {"module", "workitem", "project", "cycle", "label", "state", "user"}
for key, value in tool_args.items():
# Special-case: workspace scope sentinel should NOT trigger resolution
if key == "project_id" and isinstance(value, str) and value == "__workspace_scope__":
continue
# 1) Explicit placeholder syntax
if isinstance(value, str) and "<id of" in value:
return True
# 2) Non-UUID string passed to an *_id field for known entities
if isinstance(value, str) and key.endswith("_id"):
entity_type = key[:-3]
is_uuid = uuid_regex.match(value)
if entity_type in entity_keys and not is_uuid:
return True
# 3) List of placeholders or non-UUID id values for known entities
if isinstance(value, list):
for item in value:
# Skip sentinel in lists as well (defensive)
if key == "project_id" and isinstance(item, str) and item == "__workspace_scope__":
continue
if isinstance(item, str) and (
"<id of" in item or (key.endswith("_id") and key[:-3] in entity_keys and not uuid_regex.match(item))
):
return True
return False
async def _resolve_placeholders(self, tool_args: Dict[str, Any], tools: List, context: BatchExecutionContext) -> Optional[Dict[str, Any]]:
"""Resolve placeholders in tool arguments using retrieval tools."""
import re
resolved_args = tool_args.copy()
# Use the same UUID regex pattern as _has_placeholders for consistency
uuid_regex = re.compile(r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")
for key, value in list(resolved_args.items()):
# Special-case: workspace scope sentinel - preserve as-is
if key == "project_id" and isinstance(value, str) and value == "__workspace_scope__":
continue
if isinstance(value, str) and "<id of" in value:
match = re.search(r"<id of (\w+): ([^>]+)>", value)
if match:
entity_type, entity_name = match.groups()
# Special-case: workspace scope sentinel inside placeholder wrapper
if entity_type == "project" and entity_name.strip() == "__workspace_scope__":
resolved_args[key] = "__workspace_scope__"
continue
resolved_id = await self._resolve_entity_id(entity_type, entity_name, tools, context)
if resolved_id:
resolved_args[key] = resolved_id
else:
log.error(f"Failed to resolve {entity_type} ID for '{entity_name}'")
return None
# 2) Then auto-resolve non-UUID raw names for known *_id fields (skip placeholders)
entity_keys = {"module", "workitem", "project", "cycle", "label", "state", "user"}
for key, val in list(resolved_args.items()):
if key.endswith("_id") and key[:-3] in entity_keys and isinstance(val, str) and "<id of" not in val and not uuid_regex.match(val):
# Skip sentinel (workspace scope)
if key == "project_id" and val == "__workspace_scope__":
continue
resolved_id = await self._resolve_entity_id(key[:-3], val, tools, context)
if resolved_id:
resolved_args[key] = resolved_id
else:
log.error(f"Failed to auto-resolve {key} for '{val}'")
return None
# Now fallback to original placeholder handling for lists and stray entries
for key, value in tool_args.items():
if isinstance(value, list):
resolved_list = []
for item in value:
if isinstance(item, str) and "<id of" in item:
match = re.search(r"<id of (\w+): ([^>]+)>", item)
if match:
entity_type, entity_name = match.groups()
resolved_id = await self._resolve_entity_id(entity_type, entity_name, tools, context)
if resolved_id:
resolved_list.append(resolved_id)
else:
log.error(f"Failed to resolve {entity_type} ID for '{entity_name}'")
return None
continue
resolved_list.append(item)
resolved_args[key] = resolved_list
return resolved_args
async def _resolve_entity_id(
self, entity_type: str, entity_name: str, tools: List, context: Optional[BatchExecutionContext] = None
) -> Optional[str]:
"""Resolve a specific entity ID, first checking recent execution results, then using search tools."""
try:
# First, check if this entity was created in a previous step of the current execution
if context:
for step_id, entity_info in context.entity_results.items():
if entity_info:
# Check if this step created an entity of the requested type with the requested name
if (
entity_info.get("entity_type") == entity_type
and entity_info.get("entity_name") == entity_name
and entity_info.get("entity_id")
):
resolved_id = entity_info.get("entity_id")
return resolved_id
# If not found in execution results, use search tools
# Map entity types to their search tool names
tool_mapping = {
"module": "search_module_by_name",
"workitem": "search_workitem_by_name",
"project": "search_project_by_name",
"cycle": "search_cycle_by_name",
"label": "search_label_by_name",
"state": "search_state_by_name",
"user": "search_user_by_name",
}
tool_name = tool_mapping.get(entity_type)
if not tool_name:
log.warning(f"No tool mapping found for entity type: {entity_type}")
return None
# Find the search tool
tool_func = next((t for t in tools if getattr(t, "name", "") == tool_name), None)
if not tool_func:
log.error(f"Tool {tool_name} not found for resolving {entity_type}")
return None
# Call the search tool with the entity name
# Special case for user search which expects 'display_name' parameter
if entity_type == "user":
params = {"display_name": entity_name}
else:
params = {"name": entity_name}
if hasattr(tool_func, "ainvoke"):
result = await tool_func.ainvoke(params)
else:
result = tool_func.invoke(params)
# Parse the result to extract the ID
if isinstance(result, str):
# The tool returns a formatted string, we need to extract the ID from it
# Format: "✅ Found module 'name'\n\nResult: {'id': 'uuid', 'name': 'name', 'project_id': 'uuid'}"
try:
# Extract the Result section and parse it as a dict
if "Result:" in result:
result_section = result.split("Result:")[1].strip()
# This is a string representation of a dict, we need to parse it
import ast
result_dict = ast.literal_eval(result_section)
return result_dict.get("id")
except Exception as parse_error:
log.warning(f"Failed to parse result from {tool_name}: {parse_error}")
return None
return None
except Exception as e:
log.error(f"Error resolving {entity_type} ID for '{entity_name}': {e}")
return None
apps/pi/pi/services/chat/helpers/batch_execution_context.py
-192
View File
@@ -1,192 +0,0 @@
"""Execution context for managing batch action execution state."""
import datetime
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from uuid import UUID
from pi import logger
log = logger.getChild(__name__)
class BatchExecutionContext:
"""Manages state during batch action execution."""
def __init__(self, message_id: UUID, chat_id: UUID, user_id: UUID, workspace_id: UUID, is_project_chat: bool = False):
self.message_id = message_id
self.chat_id = chat_id
self.user_id = user_id
self.workspace_id = workspace_id
self.is_project_chat = is_project_chat
# Execution state
self.executed_actions: List[Dict[str, Any]] = []
self.current_step = 0
self.entity_results: Dict[str, Dict[str, Any]] = {} # Store created entities by step
self.failed_at_step: Optional[int] = None
self.execution_log: List[str] = []
self.start_time = datetime.datetime.utcnow()
# Status tracking
self.status = "running" # running, success, failed, partial
self.completed_count = 0
self.failed_count = 0
self.total_planned = 0
def add_execution_result(
self,
step_id: str,
tool_name: str,
result: Any,
entity_info: Optional[Dict[str, Any]] = None,
artifact_id: Optional[str] = None,
sequence: Optional[int] = None,
artifact_type: Optional[str] = None,
) -> None:
"""Record the result of an executed action."""
execution_record = {
"step": self.current_step,
"step_id": step_id,
"tool_name": tool_name,
"result": str(result),
"entity_info": entity_info,
"artifact_id": artifact_id,
"sequence": sequence,
"artifact_type": artifact_type,
"executed_at": datetime.datetime.utcnow().isoformat(),
"success": True,
}
self.executed_actions.append(execution_record)
if entity_info:
self.entity_results[f"step_{self.current_step}"] = entity_info
self.entity_results[step_id] = entity_info # Also store by step_id for easier lookup
self.completed_count += 1
self.current_step += 1
def add_execution_failure(
self,
step_id: str,
tool_name: str,
error: str,
artifact_id: Optional[str] = None,
sequence: Optional[int] = None,
artifact_type: Optional[str] = None,
) -> None:
"""Record a failed action execution."""
execution_record = {
"step": self.current_step,
"step_id": step_id,
"tool_name": tool_name,
"error": error,
"artifact_id": artifact_id,
"sequence": sequence,
"artifact_type": artifact_type,
"executed_at": datetime.datetime.utcnow().isoformat(),
"success": False,
}
self.executed_actions.append(execution_record)
self.failed_at_step = self.current_step
self.failed_count += 1
self.current_step += 1 # Increment step counter for failed actions too
# Don't immediately set status to "failed" - defer until mark_completed()
# This allows all planned actions to be attempted instead of stopping on first failure
# self.status = "failed" # REMOVED - causes early exit
log.error(f"Action execution failed: {tool_name} (step {self.current_step - 1}) - {error}")
def log_message(self, message: str) -> None:
"""Add a message to the execution log."""
timestamp = datetime.datetime.utcnow().isoformat()
log_entry = f"[{timestamp}] {message}"
self.execution_log.append(log_entry)
# log.info(message)
def get_entity_by_step(self, step_identifier: str) -> Optional[Dict[str, Any]]:
"""Get entity information by step number or step_id."""
return self.entity_results.get(step_identifier)
def get_execution_summary(self) -> Dict[str, Any]:
"""Get a summary of the execution state."""
duration = (datetime.datetime.utcnow() - self.start_time).total_seconds()
return {
"message_id": str(self.message_id),
"chat_id": str(self.chat_id),
"status": self.status,
"total_planned": self.total_planned,
"completed_count": self.completed_count,
"failed_count": self.failed_count,
"failed_at_step": self.failed_at_step,
"duration_seconds": duration,
"executed_actions": self.executed_actions,
"entity_results": self.entity_results,
"execution_log": self.execution_log[-10:] if self.execution_log else [], # Last 10 log entries
}
def get_clean_summary(self) -> Dict[str, Any]:
"""Get a clean, frontend-friendly summary without debugging information."""
duration = (datetime.datetime.utcnow() - self.start_time).total_seconds()
return {
"message_id": str(self.message_id),
"chat_id": str(self.chat_id),
"status": self.status,
"total_planned": self.total_planned,
"completed_count": self.completed_count,
"failed_count": self.failed_count,
"duration_seconds": round(duration, 2),
}
def get_clean_entity_results(self) -> Dict[str, Any]:
"""Get clean entity results without internal debugging information."""
clean_entities = {}
for key, entity_info in self.entity_results.items():
# Skip internal step keys like "step_0", "step_1"
if key.startswith("step_"):
continue
# Keep only essential entity information
if isinstance(entity_info, dict):
clean_entity = {}
for field, value in entity_info.items():
# Keep useful fields, skip internal ones
if field in ["entity_url", "entity_name", "entity_type", "entity_id"]:
clean_entity[field] = value
elif field == "tool_name":
clean_entity[field] = value
elif field == "raw_result" and len(str(value)) < 100:
# Only include short raw results
clean_entity[field] = value
if clean_entity: # Only add if we have useful information
clean_entities[key] = clean_entity
return clean_entities
def mark_completed(self) -> None:
"""Mark the batch execution as completed with proper status determination."""
if self.failed_count == 0:
self.status = "success"
elif self.completed_count == 0:
# All actions failed - complete failure
self.status = "failed"
else:
# Mixed results - some succeeded, some failed
self.status = "partial"
def can_continue(self) -> bool:
"""Check if execution can continue (not failed)."""
return self.status != "failed"
def set_total_planned(self, count: int) -> None:
"""Set the total number of planned actions."""
self.total_planned = count
apps/pi/pi/services/chat/helpers/batch_execution_helpers.py
-414
View File
@@ -1,414 +0,0 @@
"""Helper functions for batch action execution."""
import datetime
import json
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from uuid import UUID
from sqlalchemy import select
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.models.enums import ExecutionStatus
from pi.app.models.message import Message
from pi.app.models.message import MessageFlowStep
from pi.services.retrievers.pg_store.json_serializer import sanitize_execution_data
log = logger.getChild(__name__)
async def get_planned_actions_for_execution(message_id: UUID, chat_id: UUID, db: AsyncSession) -> List[Dict[str, Any]]:
"""Retrieve all planned actions for a message that are ready for execution."""
try:
# Get flow steps with planned actions - filter for TOOL steps that are planned
stmt = (
select(MessageFlowStep)
.where(MessageFlowStep.message_id == message_id) # type: ignore[arg-type]
.where(MessageFlowStep.chat_id == chat_id) # type: ignore[arg-type]
.where(MessageFlowStep.is_planned == True) # type: ignore[arg-type] # noqa: E712
.where(MessageFlowStep.is_executed == False) # type: ignore[arg-type] # noqa: E712
.order_by(MessageFlowStep.step_order) # type: ignore[arg-type]
)
result = await db.execute(stmt)
flow_steps = result.scalars().all()
planned_actions = []
for step in flow_steps:
execution_data = step.execution_data or {}
planned_action = {
"step_id": str(step.id),
"step_order": step.step_order,
"tool_name": step.tool_name,
"args": execution_data.get("args", {}),
"action_summary": execution_data.get("action_summary", {}),
"tool_id": execution_data.get("tool_id"),
"artifact_id": execution_data.get("artifact_id"), # NEW: Include artifact ID
"sequence": step.step_order, # NEW: Include planned sequence to propagate
}
planned_actions.append(planned_action)
return planned_actions
except Exception as e:
log.error(f"Error retrieving planned actions: {e}")
return []
async def get_original_user_query(message_id: UUID, db: AsyncSession) -> Optional[str]:
"""Get the original user query for the message."""
try:
stmt = select(Message).where(Message.id == message_id) # type: ignore[arg-type]
result = await db.execute(stmt)
message = result.scalar_one_or_none()
if message:
return message.content
else:
log.warning(f"Message {message_id} not found")
return None
except Exception as e:
log.error(f"Error retrieving original user query: {e}")
return None
def build_execution_prompt(original_query: str, planned_actions: List[Dict[str, Any]]) -> str:
"""Build the execution prompt for the LLM."""
# Format planned actions for the prompt
actions_text = ""
for i, action in enumerate(planned_actions, 1):
action_summary = action.get("action_summary", {})
tool_name = action.get("tool_name", "")
args = action.get("args", {})
# Create a readable description of the action
action_desc = action_summary.get("action", tool_name)
details = action_summary.get("details", "")
if details:
action_desc += f" - {details}"
actions_text += f"{i}. {action_desc}\n"
actions_text += f" Tool: {tool_name}\n"
actions_text += f" Arguments: {json.dumps(args, indent=4)}\n\n"
execution_prompt = f"""You are executing approved actions for this user request: "{original_query}"
PLANNED ACTIONS (approved by user):
{actions_text}
EXECUTION MODE: You must now EXECUTE these actions in the correct order using the available tools.
EXECUTION GUIDELINES:
- Execute each planned action exactly once using the provided tool and arguments
- Use outputs from earlier actions as inputs to later actions when needed
- If an action creates an entity (like an issue, cycle, etc.), use the returned ID in subsequent actions that need it
- Execute actions in logical dependency order (e.g., create issue first, then add to cycle)
- If an action fails, stop execution and report the specific error
- Provide clear status updates after each successful action
**CRITICAL: PLACEHOLDER HANDLING**
- Some arguments may contain placeholders that need to be resolved
- The system will automatically resolve these placeholders before tool execution
- You should execute the tools with the arguments as provided in the planned actions
- Do not attempt to manually resolve or modify the arguments
CRITICAL: You must use the EXACT tool names provided above. Do not modify, shorten, or change the tool names in any way.
AVAILABLE TOOLS: You have access to the following tool categories: {", ".join(set(action.get("tool_name", "").split("_")[0] for action in planned_actions if action.get("tool_name") and isinstance(action.get("tool_name"), str) and "_" in action.get("tool_name", "")))}
EXAMPLE: If the planned action shows tool_name: "workitems_create", you must call "workitems_create" exactly, not "create_workitem" or "workitem_create".
IMPORTANT: Start executing the actions now. Use the exact tool names and arguments provided in the planned actions above.""" # noqa: E501
return execution_prompt
def extract_entity_info(result: str, tool_name: str) -> Optional[Dict[str, Any] | List[Dict[str, Any]]]:
"""Extract entity information from tool execution result.
Returns either a single entity_info dict or a list of entity_info dicts for bulk operations.
"""
try:
result_str = str(result)
entity_info = {}
# Method 1: Check if the result is already structured JSON containing entity information
try:
if result_str.startswith("{") and result_str.endswith("}"):
result_json = json.loads(result_str)
if isinstance(result_json, dict) and "entity" in result_json:
# Extract entity information from structured response
entity_data = result_json["entity"]
if isinstance(entity_data, dict):
return entity_data
except (json.JSONDecodeError, KeyError):
pass
# Method 2: Check if the result contains entity URL information (legacy text format)
if "Entity URL:" in result_str:
lines = [ln.strip() for ln in result_str.split("\n")]
# Fast path: multiple URL lines → build list of entities
url_lines = [ln.replace("Entity URL:", "").strip() for ln in lines if ln.startswith("Entity URL:")]
if len(url_lines) > 1:
name_lines = [ln.replace("Entity Name:", "").strip() for ln in lines if ln.startswith("Entity Name:")]
type_lines = [ln.replace("Entity Type:", "").strip() for ln in lines if ln.startswith("Entity Type:")]
id_lines = [ln.replace("Entity ID:", "").strip() for ln in lines if ln.startswith("Entity ID:")]
ident_lines = [ln.replace("Entity Identifier:", "").strip() for ln in lines if ln.startswith("Entity Identifier:")]
entities: List[Dict[str, Any]] = []
max_len = len(url_lines)
for i in range(max_len):
e: Dict[str, Any] = {"entity_url": url_lines[i]}
if i < len(name_lines) and name_lines[i]:
e["entity_name"] = name_lines[i]
if i < len(type_lines) and type_lines[i]:
e["entity_type"] = type_lines[i]
if i < len(id_lines) and id_lines[i]:
e["entity_id"] = id_lines[i]
if i < len(ident_lines) and ident_lines[i]:
e["entity_identifier"] = ident_lines[i]
if e.get("entity_type") == "workitem":
e["issue_identifier"] = ident_lines[i]
entities.append(e)
# Normalize identifiers from URLs if possible
for e in entities:
try:
url = e.get("entity_url")
if isinstance(url, str) and "/browse/" in url:
after = url.split("/browse/", 1)[1]
identifier = after.split("/", 1)[0]
if "-" in identifier:
parts = identifier.split("-", 1)
if len(parts) == 2 and parts[1].isdigit():
e["issue_identifier"] = identifier
except Exception:
pass
return entities
# Single-entity path
for line in lines:
if line.startswith("Entity URL:"):
entity_info["entity_url"] = line.replace("Entity URL:", "").strip()
elif line.startswith("Entity Name:"):
entity_info["entity_name"] = line.replace("Entity Name:", "").strip()
elif line.startswith("Entity Type:"):
entity_info["entity_type"] = line.replace("Entity Type:", "").strip()
elif line.startswith("Entity ID:"):
entity_info["entity_id"] = line.replace("Entity ID:", "").strip()
elif line.startswith("Entity Identifier:"):
# Generic identifier line (project identifier or work-item unique key)
ident_val = line.replace("Entity Identifier:", "").strip()
entity_info["entity_identifier"] = ident_val
# For work-items, keep backwards-compatible key as well
try:
if entity_info.get("entity_type") == "workitem":
entity_info["issue_identifier"] = ident_val
except Exception:
pass
if entity_info:
# If we have an entity_url, attempt to parse a work-item unique key from it (e.g., WEB-821)
try:
url = entity_info.get("entity_url")
if isinstance(url, str) and "/browse/" in url:
# Expect format: .../browse/PROJECT-SEQ/
after = url.split("/browse/", 1)[1]
identifier = after.split("/", 1)[0]
# Basic validation: must contain a hyphen and end with digits
if "-" in identifier:
parts = identifier.split("-", 1)
if len(parts) == 2 and parts[1].isdigit():
entity_info["issue_identifier"] = identifier
except Exception:
pass
return entity_info
# Method 3: For tools that don't provide formatted entity info, try to infer from tool name
if tool_name and isinstance(tool_name, str) and "_create" in tool_name:
# This is a create operation, try to extract basic info
if "created successfully" in result_str.lower() or "✅" in result_str:
entity_info["tool_name"] = tool_name
entity_info["raw_result"] = result_str[:200] # Store first 200 chars
return entity_info
# Method 4: Extract entity_id from Result JSON section (fallback for when URL construction fails)
if tool_name and isinstance(tool_name, str) and "_create" in tool_name:
try:
# Look for "Result: {json_data}" pattern
if "Result: {" in result_str:
# Extract JSON from Result section
result_start = result_str.find("Result: {")
if result_start != -1:
json_start = result_start + len("Result: ")
# Find the end of the JSON object
json_str = result_str[json_start:]
# Try to parse the JSON
result_data = json.loads(json_str)
if isinstance(result_data, dict) and result_data.get("id"):
# Extract entity information from the API response
entity_info = {
"entity_id": str(result_data["id"]),
"entity_type": "epic" if "epics_create" in tool_name or "workitems_create" in tool_name else "unknown",
"tool_name": tool_name,
}
# Add entity name if available
if result_data.get("name"):
entity_info["entity_name"] = result_data["name"]
return entity_info
except (json.JSONDecodeError, ValueError, KeyError):
# JSON parsing failed, continue to fallback
pass
return None
except Exception as e:
log.warning(f"Error extracting entity info from result: {e}")
return None
async def update_flow_step_execution_status(
step_id: str,
execution_result: str,
entity_info: Optional[Dict[str, Any] | List[Dict[str, Any]]],
success: bool,
db: AsyncSession,
version_id: Optional[str] = None,
) -> None:
"""Update a flow step or artifact version to mark it as executed."""
try:
# Validate that step_id is a valid UUID
try:
step_uuid = UUID(step_id)
except ValueError as uuid_error:
log.error(f"Invalid UUID format for step_id '{step_id}': {uuid_error}")
return
# If version_id is provided, this is an edited artifact - update the ActionArtifactVersion
if version_id:
try:
from pi.services.retrievers.pg_store.action_artifact import update_action_artifact_version_execution_status
# Persist only first entity for version rows for compatibility
_single_entity: Optional[Dict[str, Any]]
if isinstance(entity_info, list):
_single_entity = entity_info[0] if entity_info else None
else:
_single_entity = entity_info
result = await update_action_artifact_version_execution_status(
db=db,
version_id=UUID(version_id),
is_executed=True,
success=success,
entity_info=_single_entity,
)
if result:
log.info(f"✅ Successfully updated ActionArtifactVersion {version_id} for step {step_id}: success={success}")
else:
log.error(f"❌ Failed to update ActionArtifactVersion {version_id} for step {step_id}")
except Exception as version_error:
log.error(f"❌ Error updating ActionArtifactVersion {version_id} for step {step_id}: {version_error}")
return
# For normal artifacts, update the MessageFlowStep
stmt = select(MessageFlowStep).where(MessageFlowStep.id == step_uuid) # type: ignore[arg-type]
query_result = await db.execute(stmt)
flow_step: Optional[MessageFlowStep] = query_result.scalars().first()
if not flow_step:
log.error(f"Flow step {step_id} not found")
return
# Mark as executed
flow_step.is_executed = True
# Set execution success status and error
if success:
flow_step.execution_success = ExecutionStatus.SUCCESS
flow_step.execution_error = None
else:
flow_step.execution_success = ExecutionStatus.FAILED
flow_step.execution_error = execution_result if "error" in execution_result.lower() else f"Execution failed: {execution_result}"
# Update execution data - create NEW dict to trigger SQLModel change detection
current_execution_data = flow_step.execution_data or {}
updated_execution_data = {
**current_execution_data, # Preserve existing data
"executed_at": str(datetime.datetime.utcnow()),
"execution_result": execution_result,
"business_success": success,
}
if entity_info:
if isinstance(entity_info, list):
updated_execution_data["entities"] = entity_info
if entity_info:
updated_execution_data["entity_info"] = entity_info[0]
else:
updated_execution_data["entity_info"] = entity_info
# Sanitize execution_data to ensure JSON serializability before saving
sanitized_execution_data = sanitize_execution_data(updated_execution_data)
# Assign NEW dictionary to trigger SQLModel change detection
flow_step.execution_data = sanitized_execution_data
# Force SQLModel to mark the field as changed
from sqlalchemy.orm import attributes
attributes.flag_modified(flow_step, "execution_data")
# Update corresponding artifact if exists (for normal flow)
artifact_id = current_execution_data.get("artifact_id")
if artifact_id:
try:
from pi.services.retrievers.pg_store.action_artifact import update_action_artifact_execution_status
# Extract single-entity info and entity_id from entity_info if available
_entity_info: Optional[Dict[str, Any]]
if isinstance(entity_info, list):
_entity_info = entity_info[0] if entity_info else None
elif isinstance(entity_info, dict):
_entity_info = entity_info
else:
_entity_info = None
entity_id = None
if isinstance(_entity_info, dict) and "entity_id" in _entity_info:
entity_id = UUID(str(_entity_info["entity_id"]))
# Update artifact execution status and persist entity_info for parity (includes entity_url)
await update_action_artifact_execution_status(
db=db,
artifact_id=UUID(artifact_id),
is_executed=True,
success=success,
entity_id=entity_id,
entity_info=_entity_info,
)
# log.info(f"Updated artifact {artifact_id} execution status: success={success}")
except Exception as artifact_error:
log.error(f"Error updating artifact {artifact_id}: {artifact_error}")
# Don't fail the main execution if artifact update fails
await db.commit()
except Exception as e:
log.error(f"Error updating flow step execution status: {e}")
await db.rollback()
# Don't re-raise - this is a tracking operation and shouldn't fail the main execution
apps/pi/pi/services/chat/helpers/build_mode_helpers.py
+1360
View File
File diff suppressed because it is too large Load Diff
apps/pi/pi/services/chat/helpers/context_manager.py
+14
View File
@@ -1,5 +1,9 @@
"""Chat context initialization and history management."""
from typing import Literal
from typing import Optional
from typing import cast
from pi import logger
log = logger.getChild(__name__)
@@ -20,8 +24,15 @@ async def initialize_chat_context(chatbot_instance, data, chat_exists, db):
# This backward compatibility code no longer does workspace resolution
is_focus_enabled = data.workspace_in_context
# Use new polymorphic fields if available, otherwise fall back to legacy fields
focus_entity_type = getattr(data, "focus_entity_type", None)
focus_entity_id = getattr(data, "focus_entity_id", None)
focus_project_id = data.project_id or None
focus_workspace_id = data.workspace_id or None
mode_raw = getattr(data, "mode", "ask") or "ask"
# Ensure mode is one of the valid literal values
mode: Optional[Literal["ask", "build"]] = cast(Literal["ask", "build"], mode_raw if mode_raw in ("ask", "build") else "ask")
if is_new:
# For new chats, the chat record should already exist from initialize-chat endpoint
@@ -49,8 +60,11 @@ async def initialize_chat_context(chatbot_instance, data, chat_exists, db):
chat_id=chat_id,
db=db,
is_focus_enabled=is_focus_enabled,
focus_entity_type=focus_entity_type,
focus_entity_id=focus_entity_id,
focus_project_id=focus_project_id,
focus_workspace_id=focus_workspace_id,
mode=mode,
)
if user_chat_preference_result["message"] != "success":
return None, "An unexpected error occurred. Please try again"
apps/pi/pi/services/chat/helpers/flow_tracking.py
+784
View File
@@ -0,0 +1,784 @@
"""Decorators for automatic flow step tracking and token management."""
import asyncio
import functools
import json
import uuid
from typing import Any
from typing import Callable
from typing import Dict
from typing import List
from typing import Optional
from pi import logger
from pi.app.models.enums import ExecutionStatus
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import MessageMetaStepType
from pi.core.db.plane_pi.lifecycle import get_streaming_db_session
from pi.services.chat.utils import standardize_flow_step_content
from pi.services.retrievers.pg_store.message import upsert_message_flow_steps
log = logger.getChild(__name__)
def track_token_usage(step_type: MessageMetaStepType):
"""
Decorator to automatically set up token tracking context for LLM calls.
Usage:
@track_token_usage(MessageMetaStepType.TOOL_ORCHESTRATION)
async def my_llm_function(self, llm_with_tools, messages, query_id, db, chat_id):
# Token tracking is automatically configured
return await llm_with_tools.ainvoke(messages)
"""
def decorator(func: Callable) -> Callable:
@functools.wraps(func)
async def wrapper(*args, **kwargs):
# Extract self, query_id, db, chat_id from args/kwargs
# Assuming standard signature: (self, llm, messages, query_id, db, chat_id, ...)
args[0] if args else None
llm = args[1] if len(args) > 1 else kwargs.get("llm")
query_id = kwargs.get("query_id") or (args[3] if len(args) > 3 else None)
db = kwargs.get("db") or (args[4] if len(args) > 4 else None)
chat_id = kwargs.get("chat_id") or (args[5] if len(args) > 5 else None)
# Set tracking context if llm supports it
if llm and hasattr(llm, "set_tracking_context"):
llm.set_tracking_context(query_id, db, step_type, chat_id=str(chat_id) if chat_id else None)
return await func(*args, **kwargs)
return wrapper
return decorator
def auto_persist_flow_steps(
step_type: FlowStepType,
tool_name: Optional[str] = None,
extract_content: Optional[Callable] = None,
extract_execution_data: Optional[Callable] = None,
timeout: float = 2.0,
on_error: str = "log", # "log", "raise", or "ignore"
):
"""
Decorator to automatically persist flow steps after function execution.
Args:
step_type: The type of flow step (TOOL, ROUTING, etc.)
tool_name: Static tool name, or None to extract from result
extract_content: Optional function to extract content from result
extract_execution_data: Optional function to extract execution_data from result
timeout: Timeout for persistence operation
on_error: How to handle persistence errors ("log", "raise", "ignore")
Usage:
@auto_persist_flow_steps(
step_type=FlowStepType.TOOL,
extract_content=lambda result: result.get('formatted_output'),
extract_execution_data=lambda result: result.get('metadata')
)
async def execute_tool(self, tool, query, query_id, chat_id, step_order, db):
# Your tool execution logic
return {"formatted_output": "...", "metadata": {...}}
"""
def decorator(func: Callable) -> Callable:
@functools.wraps(func)
async def wrapper(*args, **kwargs):
# Execute the original function
result = await func(*args, **kwargs)
# Extract required parameters from args/kwargs
# Assuming: (self, ..., query_id, chat_id, step_order, db, ...)
query_id = kwargs.get("query_id")
chat_id = kwargs.get("chat_id")
step_order = kwargs.get("step_order", 1)
db = kwargs.get("db")
if not all([query_id, chat_id, db]):
if on_error == "raise":
raise ValueError("Missing required parameters for flow step persistence")
if on_error == "log":
log.warning("Skipping flow step persistence: missing query_id, chat_id, or db")
return result
# Extract content and execution_data
content = extract_content(result) if extract_content else str(result)
execution_data = extract_execution_data(result) if extract_execution_data else {}
# Determine tool_name (static or from result)
actual_tool_name = tool_name
if not actual_tool_name and isinstance(result, dict):
actual_tool_name = result.get("tool_name")
# Build flow step
flow_step = {
"step_order": step_order,
"step_type": step_type.value,
"tool_name": actual_tool_name,
"content": standardize_flow_step_content(content, step_type),
"execution_data": execution_data,
}
# Persist asynchronously with timeout
try:
# Ensure query_id and chat_id are UUIDs
if not query_id or not chat_id:
if on_error == "raise":
raise ValueError("Missing required parameters for flow step persistence")
if on_error == "log":
log.warning("Skipping flow step persistence: missing query_id or chat_id")
return result
query_id_uuid = query_id if isinstance(query_id, uuid.UUID) else uuid.UUID(str(query_id))
chat_id_uuid = chat_id if isinstance(chat_id, uuid.UUID) else uuid.UUID(str(chat_id))
async with get_streaming_db_session() as _db:
await asyncio.wait_for(
upsert_message_flow_steps(
message_id=query_id_uuid,
chat_id=chat_id_uuid,
flow_steps=[flow_step],
db=_db,
),
timeout=timeout,
)
except asyncio.TimeoutError:
if on_error == "log":
log.warning(f"Flow step persistence timed out for {func.__name__}")
elif on_error == "raise":
raise
except Exception as e:
if on_error == "log":
log.warning(f"Flow step persistence failed for {func.__name__}: {e}")
elif on_error == "raise":
raise
return result
return wrapper
return decorator
class FlowStepCollector:
"""
Context manager for collecting flow steps during execution and persisting them in bulk.
Usage:
async with FlowStepCollector(query_id, chat_id, db) as collector:
collector.add_step(
step_order=1,
step_type=FlowStepType.TOOL,
tool_name="structured_db_tool",
content="...",
execution_data={...}
)
# Steps are automatically persisted on exit
"""
def __init__(self, query_id: uuid.UUID, chat_id: str, db: Any, timeout: float = 2.0, on_error: str = "log"):
self.query_id = query_id
self.chat_id = chat_id
self.db = db
self.timeout = timeout
self.on_error = on_error
self.steps: List[Dict[str, Any]] = []
def add_step(
self,
step_order: int,
step_type: FlowStepType,
tool_name: Optional[str] = None,
content: Any = None,
execution_data: Optional[Dict[str, Any]] = None,
**kwargs,
):
"""Add a flow step to the collection."""
step = {
"step_order": step_order,
"step_type": step_type.value if isinstance(step_type, FlowStepType) else step_type,
"tool_name": tool_name,
"content": standardize_flow_step_content(content, step_type) if content else "",
"execution_data": execution_data or {},
**kwargs,
}
self.steps.append(step)
async def __aenter__(self):
"""Enter the context manager."""
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Exit and persist all collected steps."""
if not self.steps:
return
try:
# Convert chat_id to UUID (it's typed as str)
chat_id_uuid = uuid.UUID(self.chat_id)
async with get_streaming_db_session() as _db:
await asyncio.wait_for(
upsert_message_flow_steps(
message_id=self.query_id,
chat_id=chat_id_uuid,
flow_steps=self.steps,
db=_db,
),
timeout=self.timeout,
)
except asyncio.TimeoutError:
if self.on_error == "log":
log.warning(f"Flow step collection persistence timed out ({len(self.steps)} steps)")
elif self.on_error == "raise":
raise
except Exception as e:
if self.on_error == "log":
log.warning(f"Flow step collection persistence failed: {e}")
elif self.on_error == "raise":
raise
def persist_tool_selection_steps(func: Callable) -> Callable:
"""
Decorator to persist tool selection and optional LLM reasoning steps around an LLM call.
The wrapped function must be async and accept keyword-only args:
- query_id, chat_id, db, current_step, include_context (bool),
enhanced_conversation_history (str|None), enhanced_query_for_processing (str|None), tools (list|None)
It must return an AIMessage (or compatible) which may include tool_calls and content.
Returns: (ai_message, new_current_step)
"""
@functools.wraps(func)
async def wrapper(*args, **kwargs):
query_id = kwargs.get("query_id")
chat_id = kwargs.get("chat_id")
db = kwargs.get("db")
current_step = kwargs.get("current_step", 1)
include_context = bool(kwargs.get("include_context", False))
enhanced_conversation_history = kwargs.get("enhanced_conversation_history")
enhanced_query_for_processing = kwargs.get("enhanced_query_for_processing")
tools = kwargs.get("tools") or []
ai_message = await func(*args, **kwargs)
steps: List[Dict[str, Any]] = []
try:
# Optionally store orchestration context
if (
include_context
and enhanced_conversation_history
and isinstance(enhanced_conversation_history, str)
and enhanced_conversation_history.strip()
):
steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": "tool_orchestration_context",
"content": "Context used for tool orchestration",
"execution_data": {"enhanced_conversation_history": enhanced_conversation_history},
"is_planned": False,
"is_executed": False,
})
current_step += 1
# Persist tool selection
tool_calls = getattr(ai_message, "tool_calls", None)
if tool_calls:
selected_tool_calls = [{"name": tc.get("name"), "args": tc.get("args", {}), "id": tc.get("id", "")} for tc in tool_calls]
steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": "tool_selection",
"content": standardize_flow_step_content({"selected_tools": selected_tool_calls}, FlowStepType.TOOL),
"execution_data": {
"selected_tools": selected_tool_calls,
"available_tools": [getattr(t, "name", "") for t in tools],
"query": enhanced_query_for_processing,
},
"is_planned": False,
"is_executed": False,
})
current_step += 1
# Persist LLM reasoning content if present
reasoning_text = str(getattr(ai_message, "content", "") or "").strip()
if reasoning_text:
steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": "llm_reasoning",
"content": reasoning_text,
"execution_data": {"reasoning": reasoning_text},
"is_planned": False,
"is_executed": False,
})
current_step += 1
if steps and query_id and chat_id and db:
async with get_streaming_db_session() as _db:
await upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=steps,
db=_db,
)
except Exception as e:
log.warning(f"Failed to persist tool selection/llm steps: {e}")
return ai_message, current_step
return wrapper
def persist_tool_execution_step(func: Callable) -> Callable:
"""
Decorator to persist a tool execution step (success or error).
The wrapped function must be async and accept keyword-only args:
- query_id, chat_id, db, current_step, extract_facts_fn (callable|None)
It must return the tool_result. The decorator will persist the step and return (tool_result, new_current_step).
"""
@functools.wraps(func)
async def wrapper(*args, **kwargs):
query_id = kwargs.get("query_id")
chat_id = kwargs.get("chat_id")
db = kwargs.get("db")
current_step = kwargs.get("current_step", 1)
extract_facts_fn = kwargs.get("extract_facts_fn")
# Expect first two positional args to include tool_name and tool_args via signature
# Safer: rely on kwargs provided by caller for tool_name/tool_args
tool_name = kwargs.get("tool_name") or (args[1] if len(args) > 1 else "")
tool_args = kwargs.get("tool_args") or (args[2] if len(args) > 2 else {})
try:
result = await func(*args, **kwargs)
facts: Dict[str, Any] = {}
if callable(extract_facts_fn):
try:
facts = extract_facts_fn(tool_name, result) or {}
except Exception:
facts = {}
# Normalize execution result for persistence (used by chat history reconstruction)
# Prefer structured JSON when possible, otherwise fall back to string
execution_result: str = ""
structured_result: Dict[str, Any] | List[Any] | None = None
try:
if isinstance(result, (dict, list)):
structured_result = result
# Keep a string copy to assist display formatters that expect strings
import json as _json
execution_result = _json.dumps(result, default=str)
else:
execution_result = str(result)
except Exception:
execution_result = str(result)
structured_result = None
execution_data: Dict[str, Any] = {
"tool_args": tool_args,
"tool_name": tool_name,
"execution_status": "success",
}
if facts:
execution_data["facts"] = facts
# Persist result in execution_data so retrieval can reconstruct external reasoning
if execution_result:
execution_data["execution_result"] = execution_result
if structured_result is not None:
execution_data["structured_result"] = structured_result
step = {
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": tool_name,
"content": standardize_flow_step_content(execution_data, FlowStepType.TOOL),
"execution_data": execution_data,
"is_planned": False,
"is_executed": True,
# execution_success is persisted via upsert_message_flow_steps when provided.
# We set it to SUCCESS explicitly to avoid PENDING default on retrieval.
"execution_success": ExecutionStatus.SUCCESS.value,
}
if query_id and chat_id and db:
async with get_streaming_db_session() as _db:
await upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=[step],
db=_db,
)
return result, current_step + 1
except Exception as e:
# Persist error step
error_step = {
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": tool_name,
"content": standardize_flow_step_content(
{
"tool_args": tool_args,
"execution_status": "error",
"execution_error": str(e),
},
FlowStepType.TOOL,
),
"execution_data": {
"tool_args": tool_args,
"execution_status": "error",
"execution_error": str(e),
},
"is_planned": False,
"is_executed": True,
"execution_status": ExecutionStatus.FAILED.value,
}
try:
if query_id and chat_id and db:
async with get_streaming_db_session() as _db:
await upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=[error_step],
db=_db,
)
except Exception as pe:
log.warning(f"Failed to persist error step: {pe}")
raise
return wrapper
# ------------------------------
# Pre-decorated helper functions
# ------------------------------
async def orchestrate_llm_step(llm_with_tools, messages, **kwargs):
"""Invoke LLM with tools and persist selection/reasoning using decorator."""
query_id = kwargs.get("query_id")
chat_id = kwargs.get("chat_id")
db = kwargs.get("db")
current_step = kwargs.get("current_step", 1)
include_context = bool(kwargs.get("include_context", False))
enhanced_conversation_history = kwargs.get("enhanced_conversation_history")
enhanced_query_for_processing = kwargs.get("enhanced_query_for_processing")
tools = kwargs.get("tools") or []
# Re-sync current_step with DB to avoid collisions with independently persisted steps
try:
if query_id and chat_id and db:
async with get_streaming_db_session() as _db:
from sqlalchemy import func # type: ignore
from sqlalchemy import select # type: ignore
from pi.app.models import MessageFlowStep # lazy import to avoid cycles
stmt: Any = select(func.max(MessageFlowStep.step_order)).where( # type: ignore[arg-type]
MessageFlowStep.message_id == query_id
)
result = await _db.execute(stmt)
max_step = result.scalar_one_or_none()
next_step = (max_step or 0) + 1
if current_step < next_step:
current_step = next_step
except Exception:
# best-effort; proceed with provided current_step on error
pass
# Execute LLM call
ai_message = await llm_with_tools.ainvoke(messages)
steps: List[Dict[str, Any]] = []
try:
# Optionally store orchestration context
if (
include_context
and enhanced_conversation_history
and isinstance(enhanced_conversation_history, str)
and enhanced_conversation_history.strip()
):
steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": "tool_orchestration_context",
"content": "Context used for tool orchestration",
"execution_data": {"enhanced_conversation_history": enhanced_conversation_history},
"is_planned": False,
"is_executed": False,
})
current_step += 1
log.info("orchestrate_llm_step: Added tool_orchestration_context step")
# Persist tool selection
tool_calls = getattr(ai_message, "tool_calls", None)
if tool_calls:
selected_tool_calls = [{"name": tc.get("name"), "args": tc.get("args", {}), "id": tc.get("id", "")} for tc in tool_calls]
steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": "tool_selection",
"content": standardize_flow_step_content({"selected_tools": selected_tool_calls}, FlowStepType.TOOL),
"execution_data": {
"selected_tools": selected_tool_calls,
"available_tools": [getattr(t, "name", "") for t in tools],
"query": enhanced_query_for_processing,
},
"is_planned": False,
"is_executed": False,
})
current_step += 1
log.info(f"orchestrate_llm_step: Added tool_selection step with {len(tool_calls)} tools")
# Persist LLM reasoning content if present
reasoning_text = str(getattr(ai_message, "content", "") or "").strip()
if reasoning_text:
steps.append({
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": "llm_reasoning",
"content": reasoning_text,
"execution_data": {"reasoning": reasoning_text},
"is_planned": False,
"is_executed": False,
})
current_step += 1
log.info(f"orchestrate_llm_step: Added llm_reasoning step (content length: {len(reasoning_text)})")
log.info(f"orchestrate_llm_step: Total steps to persist: {len(steps)}, query_id={query_id}, chat_id={chat_id}, db={db is not None}")
if steps and query_id and chat_id and db:
log.info(f"orchestrate_llm_step: Persisting {len(steps)} steps...")
async with get_streaming_db_session() as _db:
result = await upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id,
flow_steps=steps,
db=_db,
)
log.info(f"orchestrate_llm_step: Persistence result: {result.get("message") if isinstance(result, dict) else result}")
else:
log.warning(
f"orchestrate_llm_step: Skipping persistence - steps={len(steps)}, has_query_id={query_id is not None}, has_chat_id={chat_id is not None}, has_db={db is not None}" # noqa: E501
)
except Exception as e:
log.warning(f"Failed to persist tool selection/llm steps: {e}", exc_info=True)
return ai_message, current_step
@persist_tool_execution_step
async def execute_tool_step(tool_to_execute, tool_name: str, tool_args: Dict[str, Any], **kwargs):
"""Execute a single tool and persist success/error using decorator."""
return await tool_to_execute.ainvoke(tool_args)
async def check_and_build_clarification(*, tools: List[Any], tool_name: str, tool_args: Dict[str, Any]) -> tuple[Dict[str, Any], Any]:
"""Execute ask_for_clarification tool and return (payload, raw_result)."""
# Find tool
tool_to_execute = next((t for t in tools if getattr(t, "name", None) == tool_name), None)
try:
result = await tool_to_execute.ainvoke(tool_args) if tool_to_execute else "{}"
except Exception:
result = "{}"
# Parse payload
payload: Dict[str, Any] = {}
try:
parsed = json.loads(result) if isinstance(result, str) else result
if isinstance(parsed, dict):
payload = parsed
else:
payload = {"raw_result": str(result)}
except Exception:
payload = {"reason": tool_args.get("reason", ""), "raw_result": str(result)}
return payload, result
async def store_and_format_clarification(
*,
query_id: uuid.UUID,
chat_id: str,
db: Any,
current_step: int,
tool_name: str,
tool_args: Dict[str, Any],
result: Any,
clarification_payload: Dict[str, Any],
enhanced_query_for_processing: str | None,
tools: List[Any],
) -> tuple[str, int]:
"""Persist flow step + message clarification row and return formatted text with next step."""
# Persist flow step (pending)
next_step = await record_clarification_step(
query_id=query_id,
chat_id=chat_id,
db=db,
current_step=current_step,
tool_name=tool_name,
tool_args=tool_args,
result=result,
)
# Persist clarification record for deterministic follow-up handling
try:
from pi.services.retrievers.pg_store.clarifications import create_clarification
clar_kind = "retrieval"
bound_names = [getattr(t, "name", "") for t in tools] if tools else []
categories: list[str] = []
# Ensure original_query is not None
original_query = enhanced_query_for_processing or ""
await create_clarification(
db=db,
chat_id=uuid.UUID(str(chat_id)),
message_id=uuid.UUID(str(query_id)),
kind=clar_kind,
original_query=original_query,
payload=clarification_payload or {},
categories=categories,
method_tool_names=[],
bound_tool_names=bound_names,
)
except Exception:
# Non-fatal
pass
# Format clarification text
try:
from .tool_utils import format_clarification_as_text
formatted_text = format_clarification_as_text(clarification_payload)
except Exception:
formatted_text = json.dumps(clarification_payload, default=str)
return formatted_text, next_step
async def record_clarification_step(
*,
query_id: uuid.UUID,
chat_id: str,
db: Any,
current_step: int,
tool_name: str,
tool_args: Dict[str, Any],
result: Any,
) -> int:
"""Persist a clarification flow step as pending and return the next step index."""
try:
step = {
"step_order": current_step,
"step_type": FlowStepType.TOOL.value,
"tool_name": tool_name,
"content": standardize_flow_step_content(
{
"tool_args": tool_args,
"clarification_result": result,
},
FlowStepType.TOOL,
),
"execution_data": {
"tool_args": tool_args,
"clarification_result": result,
},
"is_planned": False,
"is_executed": False,
"execution_success": ExecutionStatus.PENDING,
}
# Convert chat_id to UUID (it's typed as str)
chat_id_uuid = uuid.UUID(chat_id)
async with get_streaming_db_session() as _db:
await upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id_uuid,
flow_steps=[step],
db=_db,
)
return current_step + 1
except Exception as e:
log.warning(f"Failed to persist clarification step: {e}")
return current_step
# ------------------------------
# Build-mode planning helpers
# ------------------------------
def build_planned_action_step(
*,
step_order: int,
tool_name: str,
action_summary: Dict[str, Any],
execution_data: Dict[str, Any],
) -> Dict[str, Any]:
"""Build a standardized planned action flow step dict (not persisted)."""
return {
"step_order": step_order,
"step_type": FlowStepType.TOOL,
"tool_name": tool_name,
"content": standardize_flow_step_content(action_summary, FlowStepType.TOOL),
"execution_data": execution_data,
"is_executed": False,
"is_planned": True,
"execution_success": ExecutionStatus.PENDING,
}
async def persist_planned_action_step(
*,
query_id: uuid.UUID,
chat_id: str,
db: Any,
step_order: int,
tool_name: str,
action_summary: Dict[str, Any],
execution_data: Dict[str, Any],
timeout: float = 2.0,
) -> Dict[str, Any]:
"""Build and persist a single planned action flow step; return the step dict.
This is useful when you want immediate durability for a planned action rather than
deferring to a batched collector.
"""
step = build_planned_action_step(
step_order=step_order,
tool_name=tool_name,
action_summary=action_summary,
execution_data=execution_data,
)
try:
chat_id_uuid = uuid.UUID(str(chat_id))
async with get_streaming_db_session() as _db:
await asyncio.wait_for(
upsert_message_flow_steps(
message_id=query_id,
chat_id=chat_id_uuid,
flow_steps=[step],
db=_db,
),
timeout=timeout,
)
except asyncio.TimeoutError:
log.warning("persist_planned_action_step: persistence timed out")
except Exception as e:
log.warning(f"persist_planned_action_step: persistence failed: {e}")
return step
apps/pi/pi/services/chat/helpers/placeholder_orchestrator.py
+546
View File
@@ -0,0 +1,546 @@
"""
Hybrid Placeholder Orchestration System
This module provides programmatic orchestration of action execution with dependency resolution.
Instead of asking the LLM to handle everything, we use:
- Python for dependency tracking and execution flow control
- Focused LLM calls for extracting specific values from tool results
"""
import json
import re
from datetime import datetime
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from pydantic import BaseModel
from pydantic import Field
from pi import logger
log = logger.getChild(__name__)
# Pydantic models for structured extraction
class ExtractedEntity(BaseModel):
"""Extracted entity information from execution context."""
entity_type: str = Field(..., description="Type of entity (e.g., 'project', 'cycle')")
entity_name: str = Field(..., description="Name of the entity as it appears in context")
entity_id: str = Field(..., description="UUID of the entity (use for *_id fields)")
entity_identifier: Optional[str] = Field(None, description="Short identifier/slug (e.g., 'PROJ', use for project identifier fields)")
class PlaceholderResolutionBatch(BaseModel):
"""Batch resolution for all placeholders in an action."""
resolutions: List[ExtractedEntity] = Field(default_factory=list, description="List of resolved entities")
class PlaceholderOrchestrator:
"""
Orchestrates execution of planned actions with smart placeholder resolution.
Key Features:
- Detects dependencies between actions
- Executes independent actions in parallel
- Resolves placeholders using focused LLM extractions
- Validates extracted values (e.g., UUID format)
"""
def __init__(self, planned_actions: List[Dict[str, Any]], chatbot, method_executor, context: Dict[str, Any], chat_id: str, message_id: str, db):
"""
Initialize the orchestrator.
Args:
planned_actions: List of actions to execute, each with tool_name and args
chatbot: PlaneChatBot instance for LLM access and tool building
method_executor: Executor for calling Plane API methods
context: Execution context (workspace_slug, user_id, etc.)
db: Database session for tracking
"""
self.planned_actions = planned_actions
self.chatbot = chatbot
self.method_executor = method_executor
self.context = context
self.db = db
self.chat_id = chat_id
self.message_id = message_id
# Execution state
self.execution_context: Dict[str, Dict[str, Any]] = {} # Stores tool results
self.results: List[Dict[str, Any]] = [] # Final execution results
self.tools_cache: Dict[str, Any] = {} # Cache of built tools
# UUID validation pattern
self.uuid_pattern = re.compile(r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")
async def execute_all(self) -> List[Dict[str, Any]]:
"""
Main orchestration loop with dependency resolution.
Returns:
List of execution results with metadata
"""
log.info(f"Starting orchestrated execution of {len(self.planned_actions)} actions")
log.info(f"\n\nPlanned actions: {self.planned_actions}\n\n")
remaining = list(self.planned_actions)
iteration = 0
max_iterations = len(self.planned_actions) * 2 # Prevent infinite loops
while remaining and iteration < max_iterations:
iteration += 1
log.info(f"Iteration {iteration}: {len(remaining)} actions remaining")
# 1. Partition actions into ready vs blocked
ready, blocked = await self._partition_by_readiness(remaining)
log.info(f"Ready: {len(ready)}, Blocked: {len(blocked)}")
# 2. Deadlock detection
if not ready and blocked:
error_msg = self._build_deadlock_error(blocked)
log.error(error_msg)
raise RuntimeError(error_msg)
# 3. Execute all ready actions
for action in ready:
try:
result = await self._execute_action(action)
self.results.append(result)
# CRITICAL: Only update context if action succeeded
if result.get("success"):
self._update_context(action, result)
else:
log.warning(f"Action {action.get("tool_name")} failed, not updating context")
remaining.remove(action)
except Exception as e:
log.error(f"Failed to execute action {action.get("tool_name")}: {e}")
# Create failure result
failure_result = {
"tool_name": action.get("tool_name"),
"result": f"❌ Failed to execute: {str(e)}",
"entity_info": None,
"artifact_id": action.get("artifact_id"),
"sequence": len(self.results) + 1,
"artifact_type": action.get("entity_type"),
"executed_at": datetime.utcnow().isoformat(),
"success": False,
"error": str(e),
}
self.results.append(failure_result)
remaining.remove(action)
# 4. Try to resolve placeholders for blocked actions
for action in blocked:
try:
await self._resolve_placeholders_in_action(action)
except Exception as e:
log.warning(f"Failed to resolve placeholders for {action.get("tool_name")}: {e}")
if iteration >= max_iterations:
log.warning(f"Reached max iterations ({max_iterations}), stopping")
log.info(f"Orchestration complete. Executed {len(self.results)} actions successfully")
return self.results
async def _partition_by_readiness(self, actions: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
"""
Split actions into ready (can execute now) vs blocked (has unresolved placeholders).
Args:
actions: List of actions to partition
Returns:
Tuple of (ready_actions, blocked_actions)
"""
ready = []
blocked = []
for action in actions:
if await self._has_unresolved_placeholders(action):
blocked.append(action)
else:
ready.append(action)
return ready, blocked
async def _has_unresolved_placeholders(self, action: Dict[str, Any]) -> bool:
"""
Check if action has placeholders that cannot yet be resolved from context.
Args:
action: Action to check
Returns:
True if action has unresolved placeholders
"""
args = action.get("args", {})
for key, value in args.items():
if self._is_placeholder(value):
# Check if we can resolve this placeholder from current context
entity_type, entity_name = self._parse_placeholder(value)
if not self._can_resolve_from_context(entity_type, entity_name):
return True
# Check lists for placeholders
if isinstance(value, list):
for item in value:
if self._is_placeholder(item):
entity_type, entity_name = self._parse_placeholder(item)
if not self._can_resolve_from_context(entity_type, entity_name):
return True
return False
async def _resolve_placeholders_in_action(self, action: Dict[str, Any]):
"""
Resolve all placeholders in an action's arguments.
Args:
action: Action with placeholders to resolve
"""
args = action.get("args", {})
for key, value in args.items():
if self._is_placeholder(value):
resolved = await self._resolve_single_placeholder(key, value)
args[key] = resolved
log.info(f"Resolved {key}: {value} -> {resolved}")
# Handle lists of placeholders
elif isinstance(value, list):
resolved_list = []
for item in value:
if self._is_placeholder(item):
resolved = await self._resolve_single_placeholder(key, item)
resolved_list.append(resolved)
else:
resolved_list.append(item)
if resolved_list != value:
args[key] = resolved_list
async def _resolve_single_placeholder(self, field_name: str, placeholder: str) -> str:
"""
Programmatically resolve placeholder by looking up entity in execution context.
NO LLM CALLS - Pure programmatic lookup for speed, reliability, and cost savings.
Args:
field_name: Name of the field (e.g., "project_id", "issues")
placeholder: Placeholder string (e.g., "<id of project: Luxury Car>")
Returns:
Resolved value for this field
Raises:
ValueError: If entity not found in context or validation fails
"""
entity_type, entity_name = self._parse_placeholder(placeholder)
if not entity_type or not entity_name:
raise ValueError(f"Failed to parse placeholder: {placeholder}")
# Build lookup key
lookup_key = f"{entity_type}:{entity_name}"
# Try exact match first
if lookup_key in self.execution_context:
context_entry = self.execution_context[lookup_key]
elif entity_name in self.execution_context:
# Fallback to name-only lookup
context_entry = self.execution_context[entity_name]
else:
raise ValueError(
f"Entity '{entity_name}' (type: {entity_type}) not found in execution context. "
f"Available entities: {list(self.execution_context.keys())}"
)
# Extract entity_info from the stored result
entity_info = context_entry.get("entity_info", {})
if not entity_info:
raise ValueError(f"No entity_info found for '{entity_name}'. " f"Stored context entry: {context_entry}")
# Programmatically select the appropriate field
# LOGIC: Almost all fields are entity references (UUID) except slugs/identifiers
if field_name in ["workspace_slug", "identifier"]:
# Slug/identifier fields → use entity_identifier (short code like "PROJ")
resolved_value = entity_info.get("entity_identifier")
if not resolved_value:
# Fallback to entity_id if identifier not present
log.warning(f"No entity_identifier for '{entity_name}', using entity_id as fallback")
resolved_value = entity_info.get("entity_id")
else:
# ALL other fields are entity references → use entity_id (UUID)
# This includes: project_id, cycle_id, issue_id, issues, assignees, members, labels, etc.
resolved_value = entity_info.get("entity_id")
if not resolved_value:
raise ValueError(f"No entity_id found for '{entity_name}' in entity_info: {entity_info}")
if not resolved_value:
raise ValueError(f"Could not extract value for field '{field_name}' from entity '{entity_name}'. " f"Entity info: {entity_info}")
# Validate the resolved value
self._validate_extracted_value(field_name, resolved_value)
log.info(f"✅ Resolved {field_name}='{placeholder}' → '{resolved_value}' " f"(from {entity_type}:{entity_name})")
return resolved_value
def _build_extraction_prompt(self, entity_type: str, entity_name: str) -> str:
"""
Build a focused prompt for extracting entity information with structured output.
Args:
entity_type: Type of entity to find (e.g., "project", "cycle")
entity_name: Name of entity to find (e.g., "Motor Bike")
Returns:
Extraction prompt string requesting structured JSON
"""
# Escape curly braces in the context JSON so LangChain doesn't treat them as variables
context_json = json.dumps(self.execution_context, indent=2)
escaped_context = context_json.replace("{", "{{").replace("}", "}}")
prompt = f"""You are extracting entity information from tool execution results.
TASK: Extract complete entity information for a specific entity
ENTITY TO FIND: {entity_type} named "{entity_name}"
CONTEXT (Tool execution results):
```json
{escaped_context}
```
INSTRUCTIONS:
You must extract ALL of the following fields for the requested entity:
1. entity_type - The type of entity (e.g., "project", "cycle", "workitem")
2. entity_name - The exact name of the entity as it appears in the context
3. entity_id - The UUID of the entity (format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx)
4. entity_identifier - The short identifier/slug (e.g., "PROJ") if available, null if not
CRITICAL RULES:
- entity_id is ALWAYS a UUID (36 characters with hyphens)
- entity_identifier is a SHORT code like "MOTO" or "PROJ" (or null if not present)
- DO NOT confuse these two! They are completely different fields!
- Look in the entity_info or entity section of the context
OUTPUT FORMAT:
Return a JSON object with this exact structure:
```json
{{{{
"entity_type": "project",
"entity_name": "Motor Bike",
"entity_id": "acf5c262-790b-4a15-bd56-501bb8863968",
"entity_identifier": "MOTO"
}}}}
```
IMPORTANT:
- Return ONLY the JSON object, no other text
- You may wrap it in markdown code blocks if needed
- All fields are required except entity_identifier (can be null)
- Make sure the entity_id is the UUID, NOT the identifier!
"""
return prompt
def _validate_extracted_value(self, field_name: str, value: str):
"""
Validate that extracted value has the correct format.
Args:
field_name: Name of the field
value: Extracted value to validate
Raises:
ValueError: If validation fails
"""
# Validate UUID fields
if field_name.endswith("_id"):
if not self._is_valid_uuid(value):
raise ValueError(
f"Field '{field_name}' requires UUID format, but got: '{value}'. " f"This looks like an identifier/slug, not a UUID."
)
# Basic presence check
if not value or value.strip() == "":
raise ValueError(f"Extracted value for '{field_name}' is empty")
def _is_valid_uuid(self, value: str) -> bool:
"""Check if value matches UUID format."""
return bool(self.uuid_pattern.match(value))
async def _execute_action(self, action: Dict[str, Any]) -> Dict[str, Any]:
"""
Execute a single action (all placeholders should be resolved at this point).
Args:
action: Action to execute
Returns:
Execution result with metadata
"""
tool_name = action["tool_name"]
args = action["args"]
log.info(f"Executing: {tool_name} with args: {args}")
# Build tool if not cached
if tool_name not in self.tools_cache:
entity_type = action.get("entity_type")
if entity_type:
tools = self.chatbot._build_method_tools(entity_type, self.method_executor, self.context)
for tool in tools:
self.tools_cache[tool.name] = tool
# Get the tool
tool = self.tools_cache.get(tool_name)
if not tool:
raise ValueError(f"Tool '{tool_name}' not found")
# Execute the tool
result = await tool.ainvoke(args)
if not isinstance(result, dict):
raise ValueError(f"Tool '{tool_name}' must return a dict, got {type(result)}")
# Build execution result
message = result.get("message") or ""
ok = bool(result.get("ok", True))
entity_info = result.get("entity")
execution_result = {
"tool_name": tool_name,
"result": message,
"entity_info": entity_info,
"artifact_id": action.get("artifact_id"),
"sequence": len(self.results) + 1,
"artifact_type": action.get("entity_type"),
"executed_at": datetime.utcnow().isoformat(),
"success": ok,
}
if not ok:
execution_result["error"] = result.get("error", "Unknown error")
return execution_result
def _update_context(self, action: Dict[str, Any], result: Dict[str, Any]):
"""
Store execution result in context for future placeholder resolution.
CRITICAL: Stores under BOTH the planned name and actual name to handle
cases where the SDK modifies the name on conflict (e.g., "Motor Bike" → "Motor Bike47119").
Args:
action: Executed action with original planned args
result: Execution result with actual entity_info
"""
entity_info = result.get("entity_info", {})
if not entity_info:
return
entity_type = entity_info.get("entity_type")
actual_name = entity_info.get("entity_name") # Name from API response (may be modified)
if not entity_type or not actual_name:
return
# Get the planned name from action args
planned_name = action.get("args", {}).get("name")
# Build all possible lookup keys
keys = [
f"{entity_type}:{actual_name}", # e.g., "project:Motor Bike47119"
actual_name, # Allow lookup by actual name alone
]
# IMPORTANT: Also store under planned name if different from actual
# This handles SDK name modifications on conflict (409)
if planned_name and planned_name != actual_name:
keys.append(f"{entity_type}:{planned_name}") # e.g., "project:Motor Bike"
keys.append(planned_name) # Allow lookup by planned name
log.info(f"⚠️ Name mismatch: planned='{planned_name}' vs actual='{actual_name}'. " f"Storing under both names for lookup.")
# Store under all keys
for key in keys:
self.execution_context[key] = result
log.info(f"✅ Stored in context: {entity_type}:{actual_name} (with {len(keys)} lookup keys)")
def _can_resolve_from_context(self, entity_type: Optional[str], entity_name: Optional[str]) -> bool:
"""
Check if we have this entity in our execution context.
Args:
entity_type: Type of entity (e.g., "project")
entity_name: Name of entity (e.g., "Motor Bike")
Returns:
True if entity is available in context
"""
if not entity_type or not entity_name:
return False
# Check if we have this entity stored
key = f"{entity_type}:{entity_name}"
return key in self.execution_context or entity_name in self.execution_context
def _is_placeholder(self, value: Any) -> bool:
"""Check if value is a placeholder string."""
return isinstance(value, str) and "<id of" in value
def _parse_placeholder(self, placeholder: str) -> Tuple[Optional[str], Optional[str]]:
"""
Parse placeholder string into entity type and name.
Args:
placeholder: Placeholder string (e.g., "<id of project: Motor Bike>")
Returns:
Tuple of (entity_type, entity_name), or (None, None) if parse fails
"""
# Pattern: '<id of {entity_type}: {entity_name}>'
match = re.match(r"<id of (\w+): (.+)>", placeholder)
if match:
return match.group(1), match.group(2).strip()
return None, None
def _build_deadlock_error(self, blocked_actions: List[Dict[str, Any]]) -> str:
"""
Build informative error message for deadlock situations.
Args:
blocked_actions: Actions that are still blocked
Returns:
Error message string
"""
placeholders = []
for action in blocked_actions:
args = action.get("args", {})
for key, value in args.items():
if self._is_placeholder(value):
placeholders.append(value)
available_entities = list(self.execution_context.keys())
return (
f"Deadlock detected: Cannot resolve placeholders.\n"
f"Blocked placeholders: {placeholders}\n"
f"Available entities in context: {available_entities}\n"
f"Possible causes:\n"
f"- Entity was not created successfully\n"
f"- Entity name mismatch between placeholder and actual name\n"
f"- Circular dependency between actions"
)
apps/pi/pi/services/chat/helpers/response_processor.py
+2 -1
View File
@@ -9,7 +9,7 @@ from pi.services.retrievers.pg_store.message import upsert_message
log = logger.getChild(__name__)
async def process_response(base_stream, chat_id, query_id, response_id, switch_llm, db, reasoning=""):
async def process_response(base_stream, chat_id, query_id, response_id, switch_llm, db, reasoning="", source=None):
"""Process streaming response and store the final result."""
final_response_ = []
@@ -35,6 +35,7 @@ async def process_response(base_stream, chat_id, query_id, response_id, switch_l
parent_id=query_id,
llm_model=switch_llm,
reasoning=reasoning,
source=source,
db=db,
)
apps/pi/pi/services/chat/helpers/tool_utils.py
+625 -148
View File
@@ -5,8 +5,10 @@ This module intentionally centralizes non-core helpers used by the action execut
so that `action_executor.execute_action_with_retrieval` stays lean and readable.
"""
import asyncio
import logging
import re
from collections.abc import AsyncIterator
from typing import Any
from typing import Dict
from typing import List
@@ -16,6 +18,86 @@ from typing import Union
log = logging.getLogger(__name__)
# build a map of tool name to category
TOOL_NAME_TO_CATEGORY_MAP: Dict[str, Dict[str, str]] = {
# Assets
"assets_create": {"entity_type": "asset", "action_type": "create", "front_facing_name": "Create Asset"},
"assets_create_user_upload": {"entity_type": "asset", "action_type": "create", "front_facing_name": "Create Asset User Upload"},
"assets_delete_user": {"entity_type": "asset", "action_type": "delete", "front_facing_name": "Delete Asset User"},
"assets_get_generic": {"entity_type": "asset", "action_type": "get", "front_facing_name": "Get Asset Generic"},
"assets_update_generic": {"entity_type": "asset", "action_type": "update", "front_facing_name": "Update Asset Generic"},
"assets_update_user": {"entity_type": "asset", "action_type": "update", "front_facing_name": "Update Asset User"},
# Attachments
"attachments_create": {"entity_type": "attachment", "action_type": "create", "front_facing_name": "Create Attachment"},
"attachments_delete": {"entity_type": "attachment", "action_type": "delete", "front_facing_name": "Delete Attachment"},
# Comments
"comments_create": {"entity_type": "comment", "action_type": "create", "front_facing_name": "Create Comment"},
"comments_delete": {"entity_type": "comment", "action_type": "delete", "front_facing_name": "Delete Comment"},
"comments_update": {"entity_type": "comment", "action_type": "update", "front_facing_name": "Update Comment"},
# Cycles
"cycles_add_work_items": {"entity_type": "cycle", "action_type": "add", "front_facing_name": "Add Work Items to Cycle"},
"cycles_archive": {"entity_type": "cycle", "action_type": "archive", "front_facing_name": "Archive Cycle"},
"cycles_create": {"entity_type": "cycle", "action_type": "create", "front_facing_name": "Create Cycle"},
"cycles_remove_work_item": {"entity_type": "cycle", "action_type": "remove", "front_facing_name": "Remove Work Item from Cycle"},
"cycles_transfer_work_items": {"entity_type": "cycle", "action_type": "transfer", "front_facing_name": "Transfer Work Items to Cycle"},
"cycles_unarchive": {"entity_type": "cycle", "action_type": "unarchive", "front_facing_name": "Unarchive Cycle"},
"cycles_update": {"entity_type": "cycle", "action_type": "update", "front_facing_name": "Update Cycle"},
# Intakes
"intake_create": {"entity_type": "intake", "action_type": "create", "front_facing_name": "Create Intake"},
"intake_delete": {"entity_type": "intake", "action_type": "delete", "front_facing_name": "Delete Intake"},
"intake_update": {"entity_type": "intake", "action_type": "update", "front_facing_name": "Update Intake"},
# Labels
"labels_create": {"entity_type": "label", "action_type": "create", "front_facing_name": "Create Label"},
"labels_update": {"entity_type": "label", "action_type": "update", "front_facing_name": "Update Label"},
# Links
"links_create": {"entity_type": "link", "action_type": "create", "front_facing_name": "Create Link"},
"links_delete": {"entity_type": "link", "action_type": "delete", "front_facing_name": "Delete Link"},
"links_update": {"entity_type": "link", "action_type": "update", "front_facing_name": "Update Link"},
# Modules
"modules_add_work_items": {"entity_type": "module", "action_type": "add", "front_facing_name": "Add Work Items to Module"},
"modules_archive": {"entity_type": "module", "action_type": "archive", "front_facing_name": "Archive Module"},
"modules_create": {"entity_type": "module", "action_type": "create", "front_facing_name": "Create Module"},
"modules_remove_work_item": {"entity_type": "module", "action_type": "remove", "front_facing_name": "Remove Work Item from Module"},
"modules_unarchive": {"entity_type": "module", "action_type": "unarchive", "front_facing_name": "Unarchive Module"},
"modules_update": {"entity_type": "module", "action_type": "update", "front_facing_name": "Update Module"},
# Pages
"pages_create_page": {"entity_type": "page", "action_type": "create", "front_facing_name": "Create Page"},
"pages_create_project_page": {"entity_type": "page", "action_type": "create", "front_facing_name": "Create Project Page"},
"pages_create_workspace_page": {"entity_type": "page", "action_type": "create", "front_facing_name": "Create Workspace Page"},
# Projects
"projects_archive": {"entity_type": "project", "action_type": "archive", "front_facing_name": "Archive Project"},
"projects_create": {"entity_type": "project", "action_type": "create", "front_facing_name": "Create Project"},
"projects_unarchive": {"entity_type": "project", "action_type": "unarchive", "front_facing_name": "Unarchive Project"},
"projects_update": {"entity_type": "project", "action_type": "update", "front_facing_name": "Update Project"},
"projects_retrieve": {"entity_type": "project", "action_type": "retrieve", "front_facing_name": "Retrieve Project"},
# Properties
"properties_create": {"entity_type": "property", "action_type": "create", "front_facing_name": "Create Property"},
"properties_create_option": {"entity_type": "property", "action_type": "create_option", "front_facing_name": "Create Property Option"},
"properties_create_value": {"entity_type": "property", "action_type": "create_value", "front_facing_name": "Create Property Value"},
"properties_delete": {"entity_type": "property", "action_type": "delete", "front_facing_name": "Delete Property"},
"properties_delete_option": {"entity_type": "property", "action_type": "delete_option", "front_facing_name": "Delete Property Option"},
"properties_update": {"entity_type": "property", "action_type": "update", "front_facing_name": "Update Property"},
"properties_update_option": {"entity_type": "property", "action_type": "update_option", "front_facing_name": "Update Property Option"},
# States
"states_create": {"entity_type": "state", "action_type": "create", "front_facing_name": "Create State"},
"states_update": {"entity_type": "state", "action_type": "update", "front_facing_name": "Update State"},
# Types
"types_create": {"entity_type": "type", "action_type": "create", "front_facing_name": "Create Type"},
"types_delete": {"entity_type": "type", "action_type": "delete", "front_facing_name": "Delete Type"},
"types_update": {"entity_type": "type", "action_type": "update", "front_facing_name": "Update Type"},
# Workitems and Epics
"create_epic": {"entity_type": "epic", "action_type": "create", "front_facing_name": "Create Epic"},
"update_epic": {"entity_type": "epic", "action_type": "update", "front_facing_name": "Update Epic"},
"workitems_create": {"entity_type": "workitem", "action_type": "create", "front_facing_name": "Create Work Item"},
"workitems_create_relation": {"entity_type": "workitem", "action_type": "update", "front_facing_name": "Create Work Item Relation"},
"workitems_update": {"entity_type": "workitem", "action_type": "update", "front_facing_name": "Update Work Item"},
"workitems_delete": {"entity_type": "workitem", "action_type": "delete", "front_facing_name": "Delete Work Item"},
# Worklogs
"worklogs_create": {"entity_type": "worklog", "action_type": "create", "front_facing_name": "Create Work Log"},
"worklogs_delete": {"entity_type": "worklog", "action_type": "delete", "front_facing_name": "Delete Work Log"},
"worklogs_update": {"entity_type": "worklog", "action_type": "update", "front_facing_name": "Update Work Log"},
}
def is_retrieval_tool(tool_name: Any) -> bool:
"""Return True if the tool is a retrieval/lookup tool.
@@ -36,29 +118,71 @@ def is_retrieval_tool(tool_name: Any) -> bool:
if name.endswith("_list") or name.endswith("_retrieve"):
return True
# Known retrieval utilities
if name in {"structured_db_tool", "vector_search_tool", "pages_search_tool", "docs_search_tool", "generic_query_tool", "fetch_cycle_details"}:
if name in {"structured_db_tool", "vector_search_tool", "pages_search_tool", "docs_search_tool", "fetch_cycle_details"}:
return True
return False
"""Tool name mapping utilities."""
"""Tool name mapping utilities and classification registry."""
# ------------------------------
# Tool classification registry
# ------------------------------
# kind: "retrieval" | "action"
# plan_only: True means never execute in planning; store as planned
TOOL_METADATA_REGISTRY: Dict[str, Dict[str, Any]] = {
# Retrieval tools (execute immediately)
"vector_search_tool": {"kind": "retrieval", "plan_only": False},
"structured_db_tool": {"kind": "retrieval", "plan_only": False},
"pages_search_tool": {"kind": "retrieval", "plan_only": False},
"docs_search_tool": {"kind": "retrieval", "plan_only": False},
"fetch_cycle_details": {"kind": "retrieval", "plan_only": False},
# Planner/system helpers
"ask_for_clarification": {"kind": "retrieval", "plan_only": False},
# Common actions (examples - fallback to heuristics for unknown tools)
"workitems_create": {"kind": "action", "plan_only": True},
"workitems_update": {"kind": "action", "plan_only": True},
"workitems_create_relation": {"kind": "action", "plan_only": True},
"modules_add_work_items": {"kind": "action", "plan_only": True},
"cycles_add_work_items": {"kind": "action", "plan_only": True},
"create_epic": {"kind": "action", "plan_only": True},
"projects_retrieve": {"kind": "retrieval", "plan_only": False},
}
def register_tool_metadata(name: str, *, kind: str, plan_only: bool = False) -> None:
"""Register or update tool metadata in the classification registry."""
if not name:
return
TOOL_METADATA_REGISTRY[name] = {"kind": kind, "plan_only": bool(plan_only)}
def get_tool_metadata(name: str) -> Dict[str, Any]:
"""Return metadata for a tool if available; empty dict if unknown."""
return TOOL_METADATA_REGISTRY.get(name, {})
def is_plan_only_tool(name: str) -> bool:
"""Return True if the tool must not be executed during planning (planned only)."""
return bool(TOOL_METADATA_REGISTRY.get(name, {}).get("plan_only", False))
import contextlib
from pi.services.schemas.chat import Agents
from pi.services.schemas.chat import RetrievalTools
def tool_name_to_agent(tool_name: str) -> str:
"""Convert tool name back to agent name for response formatting."""
tool_to_agent_map = {
"vector_search_tool": Agents.PLANE_VECTOR_SEARCH_AGENT,
"structured_db_tool": Agents.PLANE_STRUCTURED_DATABASE_AGENT,
"pages_search_tool": Agents.PLANE_PAGES_AGENT,
"docs_search_tool": Agents.PLANE_DOCS_AGENT,
"generic_query_tool": Agents.GENERIC_AGENT,
"action_executor_agent": Agents.PLANE_ACTION_EXECUTOR_AGENT,
def tool_name_to_retrieval_tool(tool_name: str) -> str:
"""Convert tool name back to retrieval tool enum value."""
tool_to_enum_map = {
"vector_search_tool": RetrievalTools.VECTOR_SEARCH_TOOL,
"structured_db_tool": RetrievalTools.STRUCTURED_DB_TOOL,
"pages_search_tool": RetrievalTools.PAGES_SEARCH_TOOL,
"docs_search_tool": RetrievalTools.DOCS_SEARCH_TOOL,
"action_executor_agent": RetrievalTools.ACTION_EXECUTOR_TOOL,
}
return tool_to_agent_map.get(tool_name, tool_name)
return tool_to_enum_map.get(tool_name, tool_name)
def tool_name_shown_to_user(tool_name: str) -> str:
@@ -66,11 +190,11 @@ def tool_name_shown_to_user(tool_name: str) -> str:
tool_to_user_map = {
"vector_search_tool": "Semantic search",
"structured_db_tool": "Database querying",
"search_current_cycle": "Find Current Cycle",
"fetch_cycle_details": "Cycle Details",
"list_recent_cycles": "Recent Cycles",
"pages_search_tool": "Semantic search of pages",
"docs_search_tool": "Semantic search of docs",
"generic_query_tool": "General Knowledge",
"action_executor_agent": "Action Execution",
# Entity search tools
"search_project_by_name": "Search Project",
@@ -92,20 +216,221 @@ def tool_name_shown_to_user(tool_name: str) -> str:
"workitems_list": "List Work-items",
"list_member_projects": "List Member Projects",
}
return tool_to_user_map.get(tool_name, tool_name)
print("received tool name: ", tool_name)
name_to_return = tool_to_user_map.get(tool_name, "")
if not name_to_return:
name_to_return = TOOL_NAME_TO_CATEGORY_MAP.get(tool_name, {}).get("front_facing_name", tool_name)
print("name to return: ", name_to_return)
return name_to_return
def agent_to_tool_name(agent_name: str) -> str:
"""Convert agent name to corresponding tool name."""
agent_to_tool_map = {
"plane_vector_search_agent": "vector_search_tool",
"plane_structured_database_agent": "structured_db_tool",
"plane_pages_agent": "pages_search_tool",
"plane_docs_agent": "docs_search_tool",
"generic_agent": "generic_query_tool",
"plane_action_executor_agent": "action_executor_agent",
def category_display_name(category: str) -> str:
"""
Convert internal category slug to a user-friendly display name.
Falls back to Title Case of the slug with underscores replaced.
"""
category = category.strip().lower()
overrides = {
"workitems": "Work items",
"worklogs": "Work logs",
"projects": "Projects",
"cycles": "Cycles",
"modules": "Modules",
"labels": "Labels",
"states": "States",
"pages": "Pages",
"assets": "Assets",
"users": "Users",
"members": "Members",
"activity": "Activity",
"attachments": "Attachments",
"comments": "Comments",
"links": "Links",
"properties": "Properties",
"types": "Types",
"intake": "Intake",
}
return agent_to_tool_map.get(agent_name, agent_name)
return overrides.get(category, category.replace("_", " ").title())
def is_uuid_like(value: Any) -> bool:
"""Check if a value looks like a UUID."""
if not isinstance(value, str):
return False
# UUID pattern: 8-4-4-4-12 hex digits with optional hyphens
uuid_pattern = re.compile(r"^[0-9a-f]{8}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{12}$", re.IGNORECASE)
return bool(uuid_pattern.match(value.strip()))
def is_url(value: Any) -> bool:
"""Check if a value looks like a URL."""
if not isinstance(value, str):
return False
return value.startswith(("http://", "https://", "ftp://"))
def clean_result_dict(data: Any, depth: int = 0) -> Any:
"""
Recursively clean a result dictionary by removing UUIDs, URLs, and technical fields.
Keeps user-friendly fields like names, identifiers, counts, dates, etc.
Args:
data: The data structure to clean (dict, list, or primitive)
depth: Current recursion depth (to prevent infinite loops)
Returns:
Cleaned data structure
"""
if depth > 10: # Safety limit
return "..."
# Fields to always exclude
EXCLUDE_FIELDS = {
"id",
"project_id",
"workspace_id",
"user_id",
"member_id",
"cycle_id",
"module_id",
"state_id",
"label_id",
"type_id",
"parent_id",
"url",
"entity_url",
"workspace",
"created_by_id",
"updated_by_id",
"owner_id",
"assignee_id",
"reporter_id",
"lead_id",
}
# Fields to always keep (even if they look like UUIDs)
KEEP_FIELDS = {
"name",
"identifier",
"title",
"description",
"count",
"total",
"status",
"priority",
"state",
"type",
"access",
"color",
"start_date",
"end_date",
"target_date",
"created_at",
"updated_at",
"is_current",
"is_active",
"is_draft",
"is_archived",
}
if isinstance(data, dict):
cleaned = {}
for key, value in data.items():
key_lower = key.lower()
# Skip excluded fields
if key_lower in EXCLUDE_FIELDS:
continue
# Skip UUID-like values unless it's a keep field
if is_uuid_like(value) and key_lower not in KEEP_FIELDS:
continue
# Skip URLs
if is_url(value):
continue
# Recursively clean nested structures
cleaned_value = clean_result_dict(value, depth + 1)
# Only include non-empty cleaned values
if cleaned_value not in (None, {}, []):
cleaned[key] = cleaned_value
return cleaned or None
elif isinstance(data, list):
cleaned_list = []
for item in data:
cleaned_item = clean_result_dict(item, depth + 1)
if cleaned_item not in (None, {}, []):
cleaned_list.append(cleaned_item)
return cleaned_list or None
else:
# Primitives: return as-is unless it's a UUID or URL
if is_uuid_like(data) or is_url(data):
return None
return data
def format_tool_message_for_display(content: Optional[str | dict]) -> str:
"""
Parse and clean tool message content for user-friendly display.
Removes UUIDs, URLs, and technical fields while preserving meaningful information.
Args:
content: Raw tool message content string (or dict that will be converted)
Returns:
Cleaned, user-friendly tool message string
"""
if not content:
return ""
# Handle case where content is still a dict (shouldn't happen but defensive)
if isinstance(content, dict):
# Extract just the message field if it's a structured response
if "message" in content:
return content["message"]
# Otherwise convert to string
content = str(content)
try:
# Split into message and result parts
parts = content.split("\n\nResult:", 1)
success_message = parts[0].strip() if parts else ""
result_part = parts[1].strip() if len(parts) > 1 else ""
if not result_part:
# No result section, just return the message as-is
return content
# We have a Result section - hide it and return just the message
# The LLM gets the full data, user only sees the clean message
return success_message or "✅ Operation completed successfully"
except Exception as e:
# On any error, return original for debugging purposes but log it
log.debug(f"Error formatting tool message: {e}")
return content
def retrieval_tool_to_tool_name(retrieval_tool: str) -> str:
"""Convert retrieval tool enum value to corresponding LangChain tool name."""
enum_to_tool_map = {
"vector_search_tool": "vector_search_tool",
"structured_db_tool": "structured_db_tool",
"pages_search_tool": "pages_search_tool",
"docs_search_tool": "docs_search_tool",
"action_executor_tool": "action_executor_agent",
}
return enum_to_tool_map.get(retrieval_tool, retrieval_tool)
# Legacy function name for backward compatibility
agent_to_tool_name = retrieval_tool_to_tool_name
def log_toolset_details(tools: List[Any], chat_id: str) -> None:
@@ -182,10 +507,37 @@ def log_toolset_details(tools: List[Any], chat_id: str) -> None:
# Action Executor helper methods
# ------------------------------
# TOOL_CALL_REASONING_INSTRUCTIONS = """**MANDATORY REASONING AND COMMUNICATION (CRITICAL - REQUIRED FOR EVERY TOOL CALL):**
# - **BEFORE EACH TOOL CALL**: You MUST explain your reasoning and intent
# - State what information you're trying to gather or what action you're planning
# - Explain why this tool is necessary for completing the user's request
# - Describe what you expect to get from the tool and how you'll use it
# - Example: "The user wants to check workitems assigned to Anil. First, I need to search for the user 'Anil' to get their ID, then I'll use that ID to filter workitems." # noqa: E501
# - **AFTER EACH TOOL CALL**: You MUST provide a brief summary of what you learned
# - Summarize key information obtained from the tool
# - Explain how this information helps with the next step
# - If the tool returned unexpected results, explain how you'll adapt
# - Example: "Found user Anil Kumar with ID xyz-123. Now I'll use this ID to search for workitems assigned to them." # noqa: E501
# - **THINKING OUT LOUD**: Express your thought process naturally
# - Share your understanding of the user's request
# - Explain your strategy for accomplishing the task
# - Mention any assumptions you're making
# - This reasoning is MANDATORY and helps with debugging and understanding your decision-making process
# - NEVER skip the reasoning - it's essential for transparency and troubleshooting""" # noqa: E501
# TOOL_CALL_REASONING_REINFORCEMENT = """**FINAL MANDATORY REQUIREMENT - REASONING FOR EVERY TOOL CALL:**
# You MUST provide clear reasoning in your response content BEFORE and AFTER each tool call:
# - BEFORE: Explain what you're about to do and why (e.g., "Let me first find user Anil's details to get their ID...") # noqa: E501
# - AFTER: Summarize what you found and your next step (e.g., "Found Anil Kumar (ID: xyz-123). Now searching for workitems assigned to them...") # noqa: E501
# - This reasoning is NOT optional - it's required for transparency and helps users understand your process
# - Even for simple searches, explain your thinking (e.g., "Searching for project 'Mobile' to get its UUID for creating workitems...")
# - Provide this in the content field surrounding your tool_calls""" # noqa: E501
# Build the planning method prompt used by the executor
def build_method_prompt(
combined_agent_query: str,
combined_tool_query: str,
project_id: Optional[str],
user_id: Optional[str],
workspace_id: Optional[str],
@@ -227,17 +579,41 @@ Use retrieval tools to gather information, then plan the modifying actions based
- Do not repeat the same action multiple times
- Do not try to execute actions - only plan them
**CRITICAL WORKFLOW:**
1. **FIRST**: Use retrieval tools to gather required information (IDs, existence checks, etc.)
2. **THEN**: Plan at least one MODIFYING ACTION that will change data
**WORKSPACE-LEVEL QUERIES (EFFICIENCY):**
|- When query is related to workitems and spans ALL projects without specifying a particular project or projects
| - DO NOT: call `list_member_projects` then query each project with `structured_db_tool` separately
| - INSTEAD: Use ONE `structured_db_tool` call WITHOUT `project_id` parameter
| - Query will automatically scope to workspace via `workspace_id` (already in context)
| - Example WRONG: `list_member_projects` + 7x `structured_db_tool` (each with different project_id)
| - Example RIGHT: Single `structured_db_tool`: "list all workitems prioritized as high"
|- For project-specific queries, include the specific `project_id`
**REASONING AND COMMUNICATION:**
- When making tool calls, ALWAYS provide a brief explanation of your reasoning in your response
- Explain why you're selecting specific tools and what you're trying to accomplish
- This helps with debugging and understanding your decision-making process
- Example: "I need to create a cycle as per the user's request, so I'll use cycles_create with the project ID."
**PROJECT FEATURES CHECK (CRITICAL - MANDATORY BEFORE CREATING PROJECT-SCOPED ENTITIES):**
- Cycles, modules, pages, workitem types, views, intake, and time tracking (worklogs) are project-level features that are enabled/disabled on a per-project basis.
- **MANDATORY WORKFLOW**: Before creating ANY of these entities (cycles_create, modules_create, pages_create_*, etc.), you MUST:
1. First get the project_id (via search_project_by_name or search_project_by_identifier if not already known)
2. **THEN** call `projects_retrieve` with that project_id to check if the feature is enabled
3. **AFTER checking**: You MUST plan the required actions using tool_calls:
- If the feature IS enabled: Plan the creation action (e.g., `cycles_create`)
- If the feature is NOT enabled: Plan BOTH `projects_update` (to enable the feature) AND the creation action (e.g., `cycles_create`)
4. **CRITICAL**: After `projects_retrieve` completes, you CANNOT stop - you MUST continue planning the modifying actions. Do NOT return only text - you MUST return tool_calls for the planned actions.
5. **REMEMBER**: `projects_retrieve` is ONLY for information gathering. The user's actual request (create cycle/module/page) is NOT complete until you plan the creation action with tool_calls. Do NOT stop after retrieval - the task is incomplete without planning the creation.
- **CRITICAL**: This check is NON-NEGOTIABLE. Never skip the `projects_retrieve` step when creating cycles, modules, pages, or other project-scoped features.
- **CRITICAL**: After retrieving project features, you MUST plan the actions - do NOT stop after retrieval. The user requested a modification (create cycle/module/page), so you MUST plan it with tool_calls.
- **EXCEPTION (NEW PROJECT IN CURRENT PLAN)**: If the target project is being CREATED in this same plan and does not yet have a real UUID, do NOT call `projects_retrieve` during planning. Use placeholders for downstream actions and defer any feature checks until after execution or when working with an existing project that has a UUID.
- Available tools:
- `projects_retrieve` tool to get details of the project features (MUST call before creating project-scoped entities)
- `projects_update` tool to update the project features (MUST include in plan if feature needs to be enabled)
- `projects_create` tool to create a new project with any/all of these features based on the user's request
**RETRIEVAL RESULT RELEVANCE (CRITICAL):**
**HARD CONSTRAINTS FOR TOOL CALLS (NON-NEGOTIABLE):**
- If the request involves modification of data like creating, updating, adding, removing, moving, assigning, archiving, or unarchiving: you MUST return one or more tool_calls for the corresponding action tools (e.g., cycles_create, workitems_create, modules_add_work_items). Returning only text is incorrect.
- For modification requests, keep text content brief - do NOT provide meta commentary like "I will plan..." or ask user to click 'Confirm'
- System handles approval automatically
- If required information is missing, call ask_for_clarification instead of skipping tool_calls.
**RETRIEVAL RESULT RELEVANCE IN PLANNING ACTIONS FOR MODIFYING REQUESTS (CRITICAL):**
- Treat retrieval results as candidates, not ground truth.
- Use results ONLY if they directly match the user's current intent and entities; otherwise ignore them.
- Never copy retrieval text verbatim into parameters like description_html unless the user explicitly asked for it.
@@ -270,26 +646,78 @@ Use retrieval tools to gather information, then plan the modifying actions based
- **No identical retries**: Do not call the same retrieval tool with the exact same parameters more than once. If it returns no/invalid results, proceed to the next fallback (within the same entity type) or ask for clarification.
- **Do not loop the same call.**
**MANDATORY PROJECT/ENTITY RESOLUTION (NO PLACEHOLDERS FOR EXISTING ENTITIES):**
- If the user mentions an EXISTING entity by NAME or IDENTIFIER, you MUST resolve it FIRST using the appropriate search tool and then use its UUID:
- project NAME → `search_project_by_name` FIRST
- project IDENTIFIER (e.g., 'HYDR', 'PARM') → `search_project_by_identifier` FIRST
- cycle/module/label/state/user/workitem NAME/IDENTIFIER → use the corresponding `search_*_by_name` / `search_*_by_identifier` tool FIRST
- Extract the UUID from the search response before using it in any subsequent tools.
- NEVER use names/identifiers directly as *_id parameters.
- NEVER emit placeholders for EXISTING entities named by the user (e.g., `project_id: "<id of project: Mattu>"` is forbidden). Resolve to a UUID instead.
- **CRITICAL**: If you need to list modules/cycles but don't have project_id, call `projects_list` FIRST to get options, then ask user which project via `ask_for_clarification` with the projects as disambiguation_options
**CRITICAL: ENTITY ID RESOLUTION AND PLACEHOLDER RULES**
**RULE 1: EXISTING ENTITIES - USE ACTUAL UUIDs (NO PLACEHOLDERS)**
When the user mentions an EXISTING entity (one that already exists in Plane):
- **YOU MUST**: Call the appropriate search tool FIRST to get its UUID
- project NAME → `search_project_by_name` → extract UUID from response
- project IDENTIFIER (e.g., 'HYDR', 'PARM') → `search_project_by_identifier` → extract UUID
- cycle/module/label/state/user/workitem NAME/IDENTIFIER → `search_*_by_name` / `search_*_by_identifier` → extract UUID
- **EXCEPTION**: User pronouns ('me', 'my', 'I', 'mine') → use User ID from USER CONTEXT directly
- **YOU MUST**: Extract the actual UUID from the search tool response (usually in the `id` field)
- **YOU MUST**: Use that extracted UUID directly in ALL subsequent tool calls
- **FORBIDDEN**: Using placeholders like `<id of workitem: ask>` for existing entities
- **FORBIDDEN**: Using names/identifiers directly as *_id parameters (e.g., `project_id: "Mobile"`)
**RULE 2: NEWLY CREATED ENTITIES - USE PLACEHOLDERS**
When planning actions that depend on entities you CREATE in the CURRENT PLAN:
- **YOU MAY**: Use placeholder references for newly created entities
- **PLACEHOLDER FORMAT**: `<id of entity_type: entity_name>`
- Module created in this plan: `<id of module: my-module>`
- Workitem created in this plan: `<id of workitem: bug fix>`
- Cycle created in this plan: `<id of cycle: Sprint 24>`
- Project created in this plan: `<id of project: my-project>`
- **EXECUTION**: The system will resolve these placeholders during execution after the entities are created
**CRITICAL EXAMPLES - MIXED EXISTING AND NEW ENTITIES:**
Example 1: "Add workitems 'Login Page' and 'Logout Page' to new cycle 'Sprint 24'"
- Step 1: search_workitem_by_name("Login Page") → Response:
{{
"id": "abc-123-uuid",
"name": "Login Page"
}}
→ Extract UUID: abc-123-uuid
- Step 2: search_workitem_by_name("Logout Page") → Response:
{{
"id": "def-456-uuid",
"name": "Logout Page"
}}
→ Extract UUID: def-456-uuid
- Step 3: Plan cycles_create(name="Sprint 24") → New cycle, will use placeholder
- Step 4: Plan cycles_add_work_items:
- cycle_id: '<id of cycle: Sprint 24>' ✅ CORRECT (newly created - use placeholder)
- issues: ['abc-123-uuid', 'def-456-uuid'] ✅ CORRECT (existing - use actual UUIDs from search)
- ❌ WRONG: issues: ['<id of workitem: Login Page>', '<id of workitem: Logout Page>'] (these are existing!)
Example 2: "Create workitem 'Fix login bug' and add it to existing module 'Backend'"
- Step 1: search_module_by_name("Backend") → Extract UUID: "module-xyz-uuid"
- Step 2: Plan workitems_create(name="Fix login bug") → New workitem, will use placeholder
- Step 3: Plan modules_add_work_items:
- module_id: 'module-xyz-uuid' ✅ CORRECT (existing - use actual UUID from search)
- issues: ['<id of workitem: Fix login bug>'] ✅ CORRECT (newly created - use placeholder)
- ❌ WRONG: module_id: '<id of module: Backend>' (this module already exists!)
**PROJECT-SCOPED ENTITIES - ADDITIONAL REQUIREMENT:**
After resolving project_id for cycles/modules/pages creation:
- **YOU MUST**: Call `projects_retrieve(project_id=...)` to verify the feature is enabled
- **EXCEPTION**: If the project itself is being CREATED in this plan, skip `projects_retrieve` (no UUID yet)
**ID VALIDATION FOR RETRIEVAL TOOLS (STRICT):**
- NEVER pass placeholders like `<id of X: name>` to `*_retrieve`, `*_list`, or `*_get` tools
- ONLY pass real UUIDs to retrieval tools
- Do NOT retrieve entities that are only PLANNED (wait until after execution when UUID exists)
- If you have a name: search first → extract UUID → then retrieve
- If you have a placeholder: do NOT call retrieval until it's resolved to a UUID
**WORKSPACE-LEVEL CONTEXT - USE PROJECT FROM HISTORY:**
- In workspace-level chats (no explicit project pre-selected), if the conversation history clearly shows a specific project selection or creation (e.g., an executed action with "Entity: <project name> (<uuid>)" or a project URL containing the UUID), you MUST include that exact UUID as `project_id` for all project-scoped tools (e.g., `workitems_create`, `workitems_update`, `modules_*`, `cycles_*`).
- Prefer the most recent project in the history when multiple appear; if multiple conflicting projects are present, disambiguate by selecting the one explicitly referenced in the current user request; otherwise ask for clarification.
- Only omit `project_id` when no project is inferable from the history or current query.
**TOOL TYPES:**
- **Retrieval tools** (search_*, *_list, *_retrieve): Execute immediately, gather info
- **Modifying actions** (*_create, *_update, *_add, *_remove): Planned for user approval
**IMPORTANT**: Only plan modifying actions if the user's request actually requires modifying data. If the request cannot be fulfilled with available tools (e.g., analytics, visualizations, external integrations), return NO_ACTIONS_PLANNED instead of creating workaround entities.
**IMPORTANT**: Only plan modifying actions if the user's request actually requires modifying data. If the request cannot be fulfilled with available tools (e.g., analytics, visualizations, external integrations), provide a polite explanation instead of creating workaround entities.
**INTERLINKED ACTIONS GUIDANCE:**
- **Multi-step operations**: When a request involves multiple related actions, you MUST plan ALL of them
@@ -305,7 +733,7 @@ Use retrieval tools to gather information, then plan the modifying actions based
**WORK-ITEM CREATION CAPABILITIES:**
**✅ CAN be set during workitems_create:**
- name, description, priority, state, assignees, labels, story_points, start_date, target_date
- **EPIC CREATION**: Use `epics_create` tool to create epics - this automatically sets the correct epic type_id
- **EPIC CREATION**: Use `create_epic` tool to create epics - this automatically sets the correct epic type_id
- **IMPORTANT**: Use workitems_create with ALL properties at once - do NOT create then update!
**WORK-ITEM RELATIONS CAPABILITIES:**
@@ -322,121 +750,56 @@ Use retrieval tools to gather information, then plan the modifying actions based
**EFFICIENCY RULE**: Always try to set as many properties as possible during creation to minimize API calls.
**SPECIFIC EXAMPLES**:
* "Make issue A block issue B" = search_workitem_by_name for A, search_workitem_by_name for B, then workitems_create_relation(issue_id=A_id, relation_type="blocking", related_issues=[B_id])
* "Mark task X as duplicate of task Y" = search for X and Y to get IDs, then workitems_create_relation(issue_id=X_id, relation_type="duplicate", related_issues=[Y_id])
* "Issue C should start after issues D and E finish" = search for C, D, E to get IDs, then workitems_create_relation(issue_id=C_id, relation_type="start_after", related_issues=[D_id, E_id])
* "Move work item to module" = search for existing work item, then modules_add_work_items
* "Create work item in module" = workitems_create, then modules_add_work_items
* "Create work item with priority and state" = workitems_create with BOTH priority AND state (single call)
* "Create work item and add to module" = workitems_create, then modules_add_work_items (two calls needed)
* "Create issue domino with Low priority and Backlog state" = workitems_create with name='domino', priority='low', state_id='backlog-uuid' (single call)
* "Move work item to cycle" = search for existing work item, then cycles_add_work_items
* Creating entities AND adding them to containers = plan BOTH create + add actions
**PLANNING DEPENDENT ACTIONS (PLACEHOLDERS ONLY FOR NEWLY CREATED ENTITIES):**
- When planning multiple actions that depend on an entity you CREATE in the same plan, you may use a descriptive reference for that newly created entity in downstream actions.
- For containers (cycles, modules, projects) or entities (work items, labels, states) that ALREADY EXIST, DO NOT use placeholders; resolve to UUIDs via search tools.
- **PLACEHOLDER SYSTEM (new entities only)**:
- Module created in this plan 'my-module' → downstream action may refer to `<id of module: my-module>`
- Workitem created in this plan 'bug fix' → downstream action may refer to `<id of workitem: bug fix>`
- Project created in this plan 'my-project' → downstream action may refer to `<id of project: my-project>`
- The execution phase will resolve these placeholders using retrieval tools ONLY for entities that were created in the current plan.
**OPTIMIZATION GUIDANCE:**
- **Descriptive references**: Use entity names/descriptions that clearly identify what should be linked
- **Eliminate redundancy**: Don't plan separate actions for each item if they can be handled in a single action if the target tool accepts a list of items
**PLACEHOLDER PLANNING EXAMPLE:**
When planning: "Create workitem 'bug fix' and add it to module 'my-module'"
- Plan: `workitems_create` with `name: 'bug fix'`
- Plan: `modules_add_work_items` with `module_id: 'my-module'` and `issues: ['bug fix']`
- The system will automatically convert 'my-module' → `<id of module: my-module>` and 'bug fix' → `<id of workitem: bug fix>`
- During execution, the LLM will resolve these placeholders using retrieval tools
**IMPORTANT**: Analyze the user's request carefully to identify ALL required actions, not just the obvious ones.
{RETRIEVAL_TOOL_DESCRIPTIONS}
**Execution Guidelines:**
- **Step 1**: Use retrieval tools (search_*_by_name, list_member_projects, structured_db_tool, etc.) to gather IDs and verify entities exist
- **Step 2**: Plan modifying actions using the gathered information
- **PROJECT LISTING**: When you need a list of active projects for scope selection or disambiguation, **PREFER** `list_member_projects` over `structured_db_tool` or `projects_list` to avoid archived/deleted projects
- **CRITICAL - USE UUIDs, NOT NAMES**: When calling tools that expect IDs (like project_id, module_id, etc.), always use the UUID from the resolved entity, NOT the entity name
- **STEP-BY-STEP PROCESS**:
1. Call search tool (e.g., `search_project_by_name`)
2. Extract the "id" field from the JSON response
3. Use that exact UUID string in subsequent tool calls
- **WRONG**: Using entity names like "MyProject", "ExampleModule", etc. for ID parameters
- **CORRECT**: Using the actual UUID string returned by search tools
- **For Relation Tasks**:
- **Step 1**: ALWAYS search for ALL work items mentioned by name or identifier to get their UUIDs
- **Step 2**: Use workitems_create_relation with the collected UUIDs
- **CRITICAL**: Never attempt relations without first collecting actual work item IDs
- Do NOT use vector_search_tool for finding projects, users, or other non-content queries
- Plan actions in logical dependency order (e.g., create issue first, then add to container)
- **Optimize similar actions**: If adding multiple items to the same target, use ONE action with a list
- **CRITICAL**: Do NOT provide workspace_slug parameters - they will be automatically provided from context
- Do not use 'search_workitem_by_identifier' if indentifier you are passing as param to it doesn't follow the format <project_identifier-workitem_sequence>, ex: 'ABCD-23'
- Brief summary:
- If you planned actions, provide a very very brief summary of the actions you planned.
- If part of the user's intent cannot be fulfilled due to API/tool limitations, provide a ONE-LINE explanatory note ONLY when strictly necessary.
- If you are unable to plan ANY modifying action after retrieval, respond with EXACTLY: NO_ACTIONS_PLANNED
- Never ask the user to confirm in text; the system UI will request confirmation.
- Use search tools for entity resolution, NOT vector_search_tool
- Do NOT provide workspace_slug - auto-provided from context
- search_workitem_by_identifier requires format: PROJECT-123
**EXAMPLE WORKFLOW FOR "MOVE" OPERATIONS:**
For "move workitem X to module Y":
1. Call `search_project_by_name` → get project UUID (executes immediately)
2. Call `search_workitem_by_name` → get work item UUID (executes immediately)
3. Call `search_module_by_name` → get module UUID (executes immediately)
4. Plan `modules_add_work_items` with the gathered UUIDs (requires user approval)
5. Use the project UUID for subsequent calls: `search_workitem_by_name` with the UUID from step 1
6. Use the project UUID for subsequent calls: `search_module_by_name` with the UUID from step 1
**MANDATORY WORKFLOW - NO EXCEPTIONS:**
1. **FIRST**: You MUST call `search_project_by_name("project-name")` to get the project UUID
2. **THEN**: Extract the UUID from the "id" field in the response
3. **FINALLY**: Use that exact UUID string in ALL subsequent tool calls
4. **FORBIDDEN**: Using project names like "Mobile", "MyProject" etc. for project_id parameters
**EXAMPLE SEQUENCE:**
- User says: "create issue in Morpheus project"
- Step 1: Call `search_project_by_name("Morpheus")` → get UUID
- Step 2: Use UUID in `workitems_create(project_id="actual-uuid-from-step-1")`
- NEVER: `workitems_create(project_id="Morpheus")` ← This will fail!
**WORKFLOW DECISION TREE:**
**CRITICAL**: "Move" means add existing entity to container - do NOT create new entities!
**IF the user's request is informational/retrieval-only (questions, searches, listing, checking status):**
1. Use retrieval/search tools to gather the requested information
2. Provide a detailed, elaborate, and neatly formatted answer based on the retrieved data. Do NOT be brief.
3. Do NOT plan any modifying actions - just return tool_calls for retrieval, then answer in your content
4. **Formatting Requirements**:
- Use "work-item" (not "issue") and "unique key" (not "Issue ID") terminology
- Suppress UUIDs - they are PII (exception: unique keys like PAI-123 are not UUIDs, can show)
- No hallucination - if no data, say so clearly without mentioning SQL/tools/internals
- Never reveal sensitive info: passwords, API keys, table names, SQL queries
- Create clickable URLs: `[PAI-123](url)` for work-items, `[name](url)` for others
- Use tables for multi-attribute data (suppress UUIDs, apply URL rules)
5. This formatting doesn't apply to modification requests - those require action planning.
**STRICT OUTPUT RULES (Very Brief Summary / NO CONFIRMATION TEXT):**
- Do NOT provide meta commentary. Do NOT say "I will plan..." or "I will update...".
- Do NOT ask the user to click 'Confirm' or otherwise prompt for approval in text. The system handles approval.
- When actions are planned: provide a very brief summary of the actions you planned.
- When no actions can be planned: use the `no_actions_planned` tool with a clear reason and optional explanation.
**STOP CONDITION:**
Note: The system uses absence of tool_calls as the signal to stop and deliver your content as the final answer. For modification requests, include at least one write tool_call in your plan.
**IF the user's request requires modifying data (create, update, delete, move, assign, etc.):**
1. Use retrieval/search tools to gather necessary information (IDs, etc.)
2. **AND** plan at least one MODIFYING ACTION (create, update, delete, add, remove, move, etc.)
3. You CANNOT stop after just searching/retrieving - you MUST plan the modifying actions
2. **AND** plan at least one MODIFYING ACTION using tool_calls (create, update, delete, add, remove, move, etc.)
3. You CANNOT stop after just searching/retrieving - you MUST plan the modifying actions with tool_calls
4. Provide a very brief summary of what you planned in your content
**IF the user's request is asking for unsupported features:**
- Analytics/visualizations (pie charts, reports, dashboards, graphs)
- External integrations (Slack, GitHub, Jira sync)
- Bulk operations (bulk delete, mass archive)
- Administrative functions (workspace settings, permissions, billing)
- File uploads/management
- Use the `no_actions_planned` tool with reason explaining why the feature is unsupported
**IF the user's request cannot be fulfilled with available tools:**
- Examples: Analytics/visualizations, external integrations, bulk operations, administrative functions, file uploads
- Use retrieval tools if relevant to understand the request
- Then provide a polite, brief explanation of why it cannot be done and what alternatives exist (if any)
- Do NOT create workaround entities (like workitems) to satisfy these requests
**CRITICAL**: Distinguish between legitimate data modification requests and requests for unsupported features. Only enforce the "must plan action" requirement for the former.
- Do NOT plan any actions - just provide the explanation in your content
""" # noqa: E501
if project_id:
method_prompt += f"\n\n**🔥 PROJECT CONTEXT (CRITICAL):**\nProject ID: {project_id}\n\n**IMPORTANT SCOPING RULES:**\n- This is a PROJECT-LEVEL chat - ALL operations are scoped to THIS PROJECT ONLY\n- When the request mentions 'current cycle', 'current module', 'work items', etc. - it means ONLY within THIS PROJECT\n- Use this project_id for ALL tools that accept project_id parameter\n- DO NOT query across all projects - scope everything to THIS specific project\n- User refers to 'this project'/'the project'/'current project' = use this project_id" # noqa: E501
else:
# Workspace-level context (no specific project)
method_prompt += f"\n\n**🌐 WORKSPACE CONTEXT (CRITICAL):**\nWorkspace ID: {workspace_id}\n\n**IMPORTANT SCOPING RULES:**\n- This is a WORKSPACE-LEVEL chat - queries can span MULTIPLE PROJECTS\n- When the request mentions 'last cycle', 'this cycle', 'work items', etc. WITHOUT specifying a project - it could be in ANY project\n- Use list_member_projects (without limit or with high limit) to get ALL projects in the workspace\n- Then iterate through projects to find relevant cycles/modules/work-items\n- CRITICAL: Do NOT limit to 1 project unless the user specifically names or refers to a specific project" # noqa: E501
method_prompt += f"\n\n**🌐 WORKSPACE CONTEXT (CRITICAL):**\nWorkspace ID: {workspace_id}\n\n**IMPORTANT SCOPING RULES:**\n- This is a WORKSPACE-LEVEL chat - queries can span MULTIPLE PROJECTS\n- When the request mentions 'last cycle', 'this cycle', 'work items', etc. WITHOUT specifying a project - it could be in ANY project\n- Use list_member_projects (with a high limit) to get ALL projects in the workspace\n- Iterate through projects ONLY for entities that are inherently project-scoped (cycles/modules/states/labels). For workspace-wide work-item queries, use a SINGLE structured_db_tool call WITHOUT project_id (it will scope via workspace_id)\n- CRITICAL: Do NOT limit to 1 project unless the user specifically names or refers to a specific project" # noqa: E501
if user_id:
method_prompt += f"\n\n**USER CONTEXT:**\nUser ID: {user_id}\nUse this when user refers to him/herself or 'I' or 'me' or 'my' or 'mine' or any other personal pronoun or any derivative of these words." # noqa: E501
@@ -444,6 +807,8 @@ For "move workitem X to module Y":
if enhanced_conversation_history and enhanced_conversation_history.strip():
method_prompt += f"\n\n**CONVERSATION HISTORY & ACTION CONTEXT:**\n{enhanced_conversation_history}\n\nBased on this conversation history, you can reference previously created entities, their IDs, and project context without needing to search for them again." # noqa: E501
log.info(f"ENHANCED_CONVERSATION_HISTORY being sent to LLM:\n{enhanced_conversation_history}")
# Inject clarification context if present (from previous turn's ask_for_clarification)
if clarification_context and isinstance(clarification_context, dict):
try:
@@ -492,8 +857,72 @@ For "move workitem X to module Y":
return method_prompt
# ------------------------------
# Clarification context builders
# ------------------------------
def build_clarification_context_block(clar_ctx: dict | None) -> str:
"""Builds a formatted clarification context block for prompts.
Expected clar_ctx keys: original_query, reason, disambiguation_options (list of dicts), answer_text
"""
try:
if not clar_ctx or not isinstance(clar_ctx, dict):
return ""
original_query_text = clar_ctx.get("original_query")
reason = clar_ctx.get("reason")
disambig_options = clar_ctx.get("disambiguation_options") or []
answer_text = clar_ctx.get("answer_text")
parts: list[str] = []
parts.append("\n\n**CLARIFICATION CONTEXT:**\n")
if original_query_text:
parts.append(f"Original user request: {original_query_text}\n")
if reason:
parts.append(f"Clarification reason: {reason}\n")
if disambig_options:
parts.append("The user was previously shown these options:\n")
for idx, opt in enumerate(disambig_options, 1):
if isinstance(opt, dict):
opt_id = opt.get("id")
opt_name = opt.get("name") or opt.get("display_name") or ""
opt_identifier = opt.get("identifier") or ""
opt_email = opt.get("email") or ""
if opt_email:
parts.append(f" {idx}. {opt_name} ({opt_email}) → UUID: {opt_id}\n")
elif opt_identifier:
parts.append(f" {idx}. {opt_name} (Identifier: {opt_identifier}) → UUID: {opt_id}\n")
else:
parts.append(f" {idx}. {opt_name} → UUID: {opt_id}\n")
if answer_text:
parts.append(f"\nUser's clarification answer: {answer_text}\n")
parts.append(
"\nIMPORTANT: The current user message is a clarification response to the original request above. "
"Use the clarification answer to resolve missing information and continue with the ORIGINAL request, "
"not as a new standalone request.\n"
)
return "".join(parts)
except Exception:
return ""
def classify_tool(tool_name: str) -> Tuple[bool, bool]:
"""Return (is_retrieval_tool, is_action_tool) based on name heuristics."""
"""Return (is_retrieval_tool, is_action_tool) using registry first, then heuristics."""
# 1) Registry check (authoritative if present)
meta = get_tool_metadata(tool_name)
if meta:
kind = str(meta.get("kind") or "").lower()
if kind == "retrieval":
return True, False
if kind == "action":
return False, True
# Fall through to heuristics if malformed
# 2) Heuristics fallback
# Include both prefix and substring patterns for robustness
read_only_patterns = [
"list_", # prefix form like list_member_projects
@@ -781,7 +1210,8 @@ def format_clarification_as_text(clarification_data: Dict[str, Any]) -> str:
REQUIRED_FIELDS_BY_TOOL: Dict[str, List[str]] = {
# Workitems
"workitems_create": ["project_id", "name"],
"workitems_update": ["issue_id", "project_id"],
# For updates, issue_id is sufficient at planning time; project_id is auto-resolved from issue_id during execution
"workitems_update": ["issue_id"],
# Modules
"modules_create": ["name", "project_id"],
"modules_add_work_items": ["module_id", "issues", "project_id"],
@@ -829,6 +1259,11 @@ def preflight_missing_required_fields(tool_name: str, tool_args: Dict[str, Any],
return missing
args = tool_args or {}
# Determine if this is an action tool; placeholders are allowed during planning for action tools
try:
_is_retrieval, _is_action = classify_tool(tool_name)
except Exception:
_is_action = False
# UUID format for strict validation of *_id fields
uuid_pattern = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.IGNORECASE)
for key in required:
@@ -844,8 +1279,11 @@ def preflight_missing_required_fields(tool_name: str, tool_args: Dict[str, Any],
if isinstance(val, str) and key == "project_id" and val == "__workspace_scope__":
is_missing = False
elif isinstance(val, str):
# Missing if placeholder or non-UUID
if "<id of" in val or not uuid_pattern.match(val):
# Placeholders like "<id of project: X>" are allowed for ACTION tools during planning
if "<id of" in val:
is_missing = not _is_action
elif not uuid_pattern.match(val):
# Non-UUID strings (e.g., names) must be resolved; still missing
is_missing = True
elif isinstance(val, dict):
# If dict provided, require a valid UUID in 'id' key
@@ -875,7 +1313,7 @@ async def handle_missing_required_fields(
chat_id: str,
tool_id: str,
current_step: int,
combined_agent_query: str,
combined_tool_query: str,
is_project_chat: Optional[bool] = None,
) -> Optional[Dict[str, Any]]:
"""Handle missing required fields by creating clarification payload with disambiguation options.
@@ -1231,7 +1669,7 @@ async def handle_missing_required_fields(
"clarification_pending": True,
"clarification_payload": clarification_payload,
# CRITICAL: Store the original query so we can reconstruct full context on clarification follow-up
"original_query": combined_agent_query,
"original_query": combined_tool_query,
},
"is_planned": False,
"is_executed": False,
@@ -1247,3 +1685,42 @@ async def handle_missing_required_fields(
except Exception:
return None
async def stream_content_in_chunks(content: str, words_per_chunk: int = 15, delay_seconds: float = 0.08) -> AsyncIterator[str]:
"""
Stream content in chunks of words with a slight delay to simulate streaming.
Args:
content: The full content to stream
words_per_chunk: Number of words to include in each chunk (default: 15)
delay_seconds: Delay between chunks in seconds (default: 0.08)
Yields:
Content chunks as strings
"""
if not content:
return
# Split by whitespace but keep delimiters to preserve formatting (newlines, tabs, etc.)
# This ensures tables and other formatted text are preserved exactly.
tokens = re.split(r"(\s+)", content)
current_chunk = []
word_count = 0
for token in tokens:
current_chunk.append(token)
# Count non-whitespace tokens as words
if token and not token.isspace():
word_count += 1
if word_count >= words_per_chunk:
yield "".join(current_chunk)
current_chunk = []
word_count = 0
await asyncio.sleep(delay_seconds)
# Yield remaining content
if current_chunk:
yield "".join(current_chunk)
apps/pi/pi/services/chat/kit.py
+100 -368
View File
@@ -5,7 +5,6 @@ from typing import Dict
from typing import List
from typing import Optional
from langchain_core.messages import BaseMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
from pydantic import UUID4
@@ -13,7 +12,6 @@ from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi import settings
from pi.agents.sql_agent import text2sql
from pi.agents.sql_agent.tools import construct_entity_urls_vectordb
from pi.agents.sql_agent.tools import extract_ids_from_sql_result
from pi.agents.sql_agent.tools import format_as_bullet_points
@@ -31,14 +29,12 @@ from pi.services.retrievers import pg_store
from pi.services.retrievers.docs_search import DocsRetriever
from pi.services.retrievers.issue_search import IssueRetriever
from pi.services.retrievers.pages_search import PageChunkRetriever
from pi.services.schemas.chat import Agents
from pi.services.schemas.chat import QueryFlowStore
from .helpers.build_mode_helpers import build_method_executor_and_context
from .mixins import AttachmentMixin
from .prompts import combination_system_prompt
from .prompts import combination_user_prompt
from .prompts import generic_prompt as GENERIC_PROMPT
from .prompts import generic_prompt_non_plane
from .prompts import title_generation_prompt
from .utils import StandardAgentResponse
from .utils import format_message_with_attachments
@@ -64,21 +60,22 @@ class ChatKit(AttachmentMixin):
use_custom_model = switch_llm is not None and not (switch_llm.startswith("gpt-5") or switch_llm == "gpt-5.1")
# Model name mapping for user-friendly names to actual LiteLLM model names
model_name_mapping = {
claude_model_name_mapping = {
"claude-sonnet-4": settings.llm_model.LITE_LLM_CLAUDE_SONNET_4,
"claude-sonnet-4-0": settings.llm_model.CLAUDE_SONNET_4_0,
"claude-sonnet-4-5": settings.llm_model.CLAUDE_SONNET_4_5,
}
tool_llm_streaming = False
if not switch_llm:
switch_llm = settings.llm_model.GPT_4_1
TOOL_LLM = settings.llm_model.GPT_4_1
tool_config = LLMConfig(model=TOOL_LLM, temperature=0.2, streaming=False)
tool_config = LLMConfig(model=TOOL_LLM, temperature=0.2, streaming=tool_llm_streaming)
else:
# Map user-friendly model names to actual model names
actual_model_name = model_name_mapping.get(switch_llm, switch_llm)
actual_model_name = claude_model_name_mapping.get(switch_llm, switch_llm)
if switch_llm in model_name_mapping:
if switch_llm in claude_model_name_mapping:
# This is a Claude model
if actual_model_name in [settings.llm_model.CLAUDE_SONNET_4_0, settings.llm_model.CLAUDE_SONNET_4_5]:
# Direct Anthropic API models
@@ -90,7 +87,7 @@ class ChatKit(AttachmentMixin):
tool_config = LLMConfig(
model=TOOL_LLM,
temperature=0.2,
streaming=False,
streaming=tool_llm_streaming,
base_url=settings.llm_config.LITE_LLM_HOST,
api_key=settings.llm_config.LITE_LLM_API_KEY,
)
@@ -102,7 +99,7 @@ class ChatKit(AttachmentMixin):
TOOL_LLM = "gpt-5" # Use base GPT-5 model name for OpenAI API
tool_config = LLMConfig(
model=TOOL_LLM,
streaming=False,
streaming=tool_llm_streaming,
reasoning_effort="medium",
use_responses_api=settings.llm_config.GPT5_USE_RESPONSES_API, # Configurable via env var
)
@@ -111,7 +108,7 @@ class ChatKit(AttachmentMixin):
TOOL_LLM = "gpt-5" # Use base GPT-5 model name for OpenAI API
tool_config = LLMConfig(
model=TOOL_LLM,
streaming=False,
streaming=tool_llm_streaming,
reasoning_effort="low",
use_responses_api=settings.llm_config.GPT5_USE_RESPONSES_API, # Configurable via env var
)
@@ -125,7 +122,7 @@ class ChatKit(AttachmentMixin):
else:
# This is a regular OpenAI model
TOOL_LLM = switch_llm
tool_config = LLMConfig(model=TOOL_LLM, temperature=0.2, streaming=False)
tool_config = LLMConfig(model=TOOL_LLM, temperature=0.2, streaming=tool_llm_streaming)
# Initialize LLMs using LLMFactory with switch_llm support
if use_custom_model:
@@ -163,15 +160,13 @@ class ChatKit(AttachmentMixin):
self.chat_llm = llms.get_chat_llm(switch_llm)
# Log key model info
chat_model = getattr(self.chat_llm, "model_name", "unknown")
reasoning_effort = "N/A"
getattr(self.chat_llm, "model_name", "unknown")
if switch_llm == "gpt-5-standard":
reasoning_effort = "medium"
pass
elif switch_llm == "gpt-5-fast":
reasoning_effort = "low"
log.info(f"Using model: {chat_model} (reasoning_effort: {reasoning_effort}), switch_llm: {switch_llm}")
pass
self.tool_llm = create_openai_llm(tool_config)
self.tool_llm = llms.LazyLLM(lambda: create_openai_llm(tool_config))
self.issue_retriever = IssueRetriever(
num_docs=settings.chat.NUM_SIMILAR_DOCS, chunk_similarity_threshold=settings.vector_db.ISSUE_VECTOR_SEARCH_CUTOFF
)
@@ -190,19 +185,19 @@ class ChatKit(AttachmentMixin):
self.vector_search_issue_ids: list[str] = []
self.vector_search_page_ids: list[str] = []
self.current_context: dict[str, Any] = {}
# Store standardized agent responses for URL access
self.agent_responses: dict[str, Dict[str, Any]] = {}
# Store standardized tool responses for URL access
self.tool_responses: dict[str, Dict[str, Any]] = {}
# Token tracking context (set externally when needed)
self._token_tracking_context: Optional[Dict[str, Any]] = None
def _store_agent_response(self, tool_name: str, response: Dict[str, Any]) -> None:
"""Store standardized agent response for later URL access"""
self.agent_responses[tool_name] = response
"""Store standardized tool response for later URL access"""
self.tool_responses[tool_name] = response
def _get_stored_response(self, tool_name: str) -> Optional[Dict[str, Any]]:
"""Get stored agent response"""
return self.agent_responses.get(tool_name)
"""Get stored tool response"""
return self.tool_responses.get(tool_name)
def set_token_tracking_context(self, message_id: UUID4, db: AsyncSession, chat_id: Optional[str] = None) -> None:
"""Set token tracking context for this chat session"""
@@ -365,13 +360,13 @@ Provide concise, relevant context from the attachment(s):"""
entity_urls = None
try:
extracted_ids = extract_ids_from_sql_result(query_execution_result)
log.info(f"ChatID: {chat_id} - _create_entity_urls_for_db_search: Extracted IDs: {extracted_ids}")
# Construct URLs using dynamic function (no need for user_meta)
if any(extracted_ids.values()):
api_base_url = settings.plane_api.FRONTEND_URL
log.info(f"ChatID: {chat_id} - kit.py:_create_entity_urls_for_db_search: API Base URL: {api_base_url}")
entity_urls = await construct_entity_urls_vectordb(entity_ids=extracted_ids, api_base_url=api_base_url)
log.info(f"ChatID: {chat_id} - kit.py:_create_entity_urls_for_db_search: Entity URLs: {entity_urls}")
# intermediate_results["entity_urls"] = entity_urls
query_flow_store["tool_response"] += f"Entity extraction: {sum(len(ids) for ids in extracted_ids.values())} entities found\n"
@@ -645,35 +640,6 @@ Provide concise, relevant context from the attachment(s):"""
# The action executor intercepts this tool call and streams a dedicated clarification event.
return json.dumps(payload)
@tool
async def no_actions_planned(
reason: str,
explanation: Optional[str] = None,
) -> str:
"""Use this when you determine that no actions need to be planned for the user's request.
This should be used when:
- The request is asking for information only (no data modification needed)
- The request is for unsupported features (analytics, external integrations, etc.)
- The request cannot be fulfilled with available tools
- All necessary information has been retrieved and no further actions are needed
Args:
reason: Short description of why no actions are needed (e.g., "Request is for information only").
explanation: Optional detailed explanation for the user about why no actions were planned.
Returns:
JSON string with the reasoning for downstream handling.
"""
payload: Dict[str, Any] = {
"reason": reason,
"explanation": explanation or "",
"token": "NO_ACTIONS_PLANNED",
}
# The action executor will intercept this and handle the no-actions flow
return json.dumps(payload)
@tool
async def vector_search_tool(query: str) -> str:
"""Search for issues using semantic vector search. Use this first when you need to find issues related to specific topics or keywords.
@@ -729,7 +695,10 @@ Provide concise, relevant context from the attachment(s):"""
self._store_agent_response("structured_db_tool", response)
# Extract and return the results text
return StandardAgentResponse.extract_results(response)
# result = StandardAgentResponse.extract_results(response)
result = StandardAgentResponse.format_response_with_entity_urls(response)
log.info(f"ChatID: {chat_id} - kit.py:structured_db_tool: format_response_with_entity_urls result: {result}")
return result
except Exception as e:
log.error(f"Error in structured_db_tool: {str(e)}")
return f"Error querying database: {str(e)}"
@@ -1135,32 +1104,21 @@ Provide concise, relevant context from the attachment(s):"""
log.error(f"Error in docs_search_tool: {str(e)}")
return f"Error searching documentation: {str(e)}"
@tool
async def generic_query_tool(query: str) -> str:
"""Handle general questions that don't require specific data retrieval from the workspace."""
try:
# Get attachment blocks from the shared context
attachment_blocks = self.get_current_attachment_blocks()
result = await self.handle_generic_query(query, user_id, conversation_history, attachment_blocks=attachment_blocks)
# Store the standardized response for consistency
self._store_agent_response("generic_query_tool", result)
# Extract and return the results text
return StandardAgentResponse.extract_results(result)
except Exception as e:
log.error(f"Error in generic_query_tool: {str(e)}")
return f"Error handling generic query: {str(e)}"
# Store fetch_cycle_details in self for later access by _get_selected_tools
self._fetch_cycle_details_tool = fetch_cycle_details
return [ask_for_clarification, vector_search_tool, structured_db_tool, pages_search_tool, docs_search_tool, generic_query_tool]
return [
ask_for_clarification,
vector_search_tool,
structured_db_tool,
pages_search_tool,
docs_search_tool,
fetch_cycle_details,
]
async def _get_selected_tools(
async def _create_tools_for_ask_mode(
self,
db,
selected_agents,
user_meta,
workspace_id,
workspace_slug,
@@ -1171,112 +1129,17 @@ Provide concise, relevant context from the attachment(s):"""
conversation_history,
message_id,
is_project_chat=None,
pi_sidebar_open=None,
sidebar_open_url=None,
workspace_in_context=True,
chatbot_instance=None,
):
"""Get tools based on selected agents, with dynamic OAuth token retrieval for action agents."""
"""Create tools for ask mode.
# Check for action executor agent
is_action_agent = any(agent_query.agent == Agents.PLANE_ACTION_EXECUTOR_AGENT for agent_query in selected_agents)
if is_action_agent:
# Build context for selector and executor - prefer provided workspace_slug; resolve only if missing
if not workspace_slug:
from pi.app.api.v1.helpers.plane_sql_queries import get_workspace_slug
workspace_slug = await get_workspace_slug(workspace_id)
if not workspace_slug:
log.warning(f"Could not resolve workspace slug for workspace {workspace_id}")
return self._create_auth_error_tools(f"Could not resolve workspace information for {workspace_id}")
# Get OAuth token for user and workspace (development mode will return hardcoded API key)
access_token = await self._get_oauth_token_for_user(db, user_id, workspace_id)
if not access_token:
# No valid OAuth token - return special "auth required" tool
log.info(f"OAuth required for user {user_id} in workspace {workspace_id}")
return self._create_auth_required_tools(
workspace_id,
user_id,
chat_id,
str(message_id),
is_project_chat=is_project_chat,
project_id=project_id,
pi_sidebar_open=pi_sidebar_open,
sidebar_open_url=sidebar_open_url,
workspace_slug=workspace_slug,
)
# Initialize the executor with OAuth token or API key
action_tools: list = []
try:
from pi.services.actions.category_selector import CategorySelector
from pi.services.actions.method_executor import MethodExecutor
from pi.services.actions.plane_actions_executor import PlaneActionsExecutor
# Create actions executor
if access_token and access_token.startswith("plane_api_"):
actions_executor = PlaneActionsExecutor(api_key=access_token, base_url=settings.plane_api.HOST)
else:
actions_executor = PlaneActionsExecutor(access_token=access_token, base_url=settings.plane_api.HOST)
# Create hierarchical tools
category_selector = CategorySelector()
method_executor = MethodExecutor(actions_executor)
@tool
async def get_available_plane_actions(user_intent: str) -> str:
"""
Return available Plane API methods across all categories (advisory catalog).
Used to inform the LLM router; does not decide categories itself.
Categories include: cycles, workitems (for issues/tasks), projects, labels, states, modules, pages, assets, users, intake, members, activity, attachments, comments, links, properties, types, worklogs.
""" # noqa: E501
categories = category_selector.get_available_categories()
lines: list[str] = []
lines.append("Available Plane action categories and methods:\n")
for cat, description in categories.items():
try:
cat_methods = method_executor.get_category_methods(cat)
method_names = ", ".join(cat_methods.keys()) if cat_methods else "-"
except Exception as exc:
log.warning(f"Failed to get methods for category '{cat}': {exc}")
method_names = "(error retrieving methods)"
lines.append(f"- {cat}: {description}")
lines.append(f" Methods: {method_names}")
return "\n".join(lines)
action_tools = [get_available_plane_actions]
except Exception as e:
log.error(f"Error creating action tools: {e}")
return self._create_auth_error_tools(str(e))
# ALSO include existing retrieval tools so LLM can mix retrieval + actions
all_static_tools = self._create_tools(
db,
user_meta,
workspace_id,
project_id,
user_id,
chat_id,
query_flow_store,
conversation_history,
message_id,
is_project_chat=is_project_chat,
)
# Combine action tools with retrieval tools
combined_tools = action_tools + all_static_tools
return combined_tools
# Default: static retrieval tools
agent_to_tool_map = {
Agents.PLANE_STRUCTURED_DATABASE_AGENT: "structured_db_tool",
Agents.PLANE_VECTOR_SEARCH_AGENT: "vector_search_tool",
Agents.PLANE_PAGES_AGENT: "pages_search_tool",
Agents.PLANE_DOCS_AGENT: "docs_search_tool",
Agents.GENERIC_AGENT: "generic_query_tool",
}
Args:
workspace_in_context: If False, filters out workspace-specific tools like
structured_db_tool, vector_search_tool, pages_search_tool, fetch_cycle_details,
entity search tools, and ask_for_clarification. Only keeps docs_search_tool.
"""
# Create all the retrieval tools except the generic query tool, which is no longer needed in ask mode
all_tools = self._create_tools(
db,
user_meta,
@@ -1289,15 +1152,19 @@ Provide concise, relevant context from the attachment(s):"""
message_id,
is_project_chat=is_project_chat,
)
selected_tool_names = [agent_to_tool_map.get(agent_query.agent) for agent_query in selected_agents]
tools = [tool for tool in all_tools if tool.name in selected_tool_names]
# Inject entity search tools and fetch_cycle_details so retrieval-only flows can disambiguate entities
# and get cycle details without SQL generation
try:
# Filter tools based on workspace_in_context
if workspace_in_context:
# Include all workspace-specific tools when workspace is in context
retrieval_tools = [
tool
for tool in all_tools
if tool.name in ["vector_search_tool", "structured_db_tool", "pages_search_tool", "docs_search_tool", "fetch_cycle_details"]
]
clarification_tool = next((t for t in all_tools if getattr(t, "name", "") == "ask_for_clarification"), None)
# Add all the entity search tools
from pi.services.actions.tools.entity_search import get_entity_search_tools
# Build minimal context for entity search and cycle details tool
search_ctx = {
"workspace_id": workspace_id,
"workspace_slug": workspace_slug,
@@ -1305,33 +1172,35 @@ Provide concise, relevant context from the attachment(s):"""
"user_id": user_id,
}
entity_tools = get_entity_search_tools(method_executor=None, context=search_ctx) or []
log.info(f"ChatID: {chat_id} - Got {len(entity_tools)} entity search tools")
retrieval_tools.extend(entity_tools)
# Add fetch_cycle_details tool (stored in self during _create_tools)
if hasattr(self, "_fetch_cycle_details_tool"):
entity_tools.append(self._fetch_cycle_details_tool)
log.info(f"ChatID: {chat_id} - Added fetch_cycle_details to multi-agent tools")
else:
log.warning(f"ChatID: {chat_id} - fetch_cycle_details tool not available (was _create_tools called?)")
# add projects_retrieve tool till the retrieval and entity search tools are merged
method_executor, context, workspace_slug = await build_method_executor_and_context(
chatbot_instance=chatbot_instance,
user_id=user_id,
workspace_id=workspace_id,
project_id=project_id,
conversation_history=conversation_history,
user_meta=user_meta,
is_project_chat=is_project_chat,
chat_id=chat_id,
db=db,
)
project_tools = get_tools_for_category("projects", method_executor=method_executor, context=context) or []
project_tools = [t for t in project_tools if getattr(t, "name", "") in ["projects_retrieve"]]
retrieval_tools.extend(project_tools)
# Merge while avoiding duplicates by name
existing = {getattr(t, "name", "") for t in tools}
for t in entity_tools:
if getattr(t, "name", "") not in existing:
tools.append(t)
log.info(f"ChatID: {chat_id} - After adding entity tools, total tools: {len(tools)}, names: {[t.name for t in tools]}")
except Exception as e:
# Best-effort only; skip if unavailable
log.error(f"ChatID: {chat_id} - Error adding entity tools: {e}")
pass
# Always include clarification tool for retrieval flows so LLM can resolve ambiguity
try:
clar_tool = next((t for t in all_tools if getattr(t, "name", "") == "ask_for_clarification"), None)
if clar_tool and all(getattr(t, "name", "") != "ask_for_clarification" for t in tools):
tools = [clar_tool] + tools
except Exception:
pass
return tools
return retrieval_tools + [clarification_tool] if clarification_tool else retrieval_tools
else:
# When workspace is not in context, only include non-workspace-specific tools
log.info(f"ChatID: {chat_id} - Workspace not in context, filtering out workspace-specific tools")
retrieval_tools = [
tool
for tool in all_tools
if tool.name in ["docs_search_tool"] # Only general documentation search
]
# Don't include ask_for_clarification as it's designed for workspace entity clarifications
return retrieval_tools
def _build_method_tools(self, category: str, method_executor, context: dict):
"""Build method-specific tools for the selected category using modular structure"""
@@ -1414,149 +1283,6 @@ Provide concise, relevant context from the attachment(s):"""
return title
@llm_error_handler(
fallback_message="I'm having trouble processing your request right now. Please try again later.", max_retries=2, log_context="[GENERIC_QUERY]"
)
async def handle_generic_query(
self,
query: str,
user_id: str,
conversation_history: list[BaseMessage],
user_meta: Optional[dict[str, Any]] = None,
attachment_blocks: Optional[List[Dict[str, Any]]] = None,
include_user_context_on_followup: bool = False,
) -> Dict[str, Any]:
"""
Handle generic queries using LLM without database or vector search.
input: query - user query
output: Generic response based on LLM knowledge only
"""
if user_meta is None:
user_meta = {}
date_time_context = await get_current_timestamp_context(user_id)
first_name = user_meta.get("first_name") or user_meta.get("firstName", "")
last_name = user_meta.get("last_name") or user_meta.get("lastName", "")
# Compose user context for possible use in system prompt
user_only_context = f"**User's Firstname**: {first_name} and Lastname: {last_name}"
# Compose message content with attachments if present
time_context_query = f"{query}\n\n**Context**: {date_time_context}"
message_content = time_context_query if not attachment_blocks else format_message_with_attachments(time_context_query, attachment_blocks)
recent_history = [(m.type, m.content) for m in conversation_history]
system_prompt = GENERIC_PROMPT
if conversation_history:
system_prompt = f"{system_prompt}\n\n{date_time_context}"
if include_user_context_on_followup:
log.debug("[GENERIC_QUERY] Adding user context on follow-up (include_user_context_on_followup=True)")
system_prompt = f"{system_prompt}\n\n{user_only_context}"
system_prompt = f"{system_prompt}\n\nSkip greetings and get straight to the point."
else:
system_prompt = f"{system_prompt}\n\n{date_time_context}\n\n{user_only_context}"
recent_history = [(m.type, m.content) for m in conversation_history]
generic_prompt = ChatPromptTemplate.from_messages([
("system", system_prompt),
("placeholder", "{recent_history}"),
("human", "{question}"),
])
# Set tracking context for generic query
if self._token_tracking_context:
self.chat_llm.set_tracking_context(
self._token_tracking_context["message_id"],
self._token_tracking_context["db"],
MessageMetaStepType.COMBINATION,
chat_id=self._token_tracking_context.get("chat_id"),
)
# Use LLM chain
generic_llm_chain = generic_prompt | self.chat_llm
llm_response = await generic_llm_chain.ainvoke({"question": message_content, "recent_history": recent_history})
response_content = llm_response.content if hasattr(llm_response, "content") else str(llm_response)
return StandardAgentResponse.create_response(response_content)
async def handle_generic_query_stream(
self,
query: str,
user_id: str,
conv_history: list[BaseMessage],
user_meta: dict[str, Any],
non_plane: bool = False,
include_user_context_on_followup: bool = False,
attachment_blocks: Optional[List[Dict[str, Any]]] = None,
) -> AsyncIterator[str]:
date_time_context = await get_current_timestamp_context(user_id)
first_name = user_meta.get("first_name") or user_meta.get("firstName", "")
last_name = user_meta.get("last_name") or user_meta.get("lastName", "")
user_only_context = f"**User's Firstname**: {first_name} and Lastname: {last_name}"
if non_plane:
system_prompt = generic_prompt_non_plane
else:
system_prompt = GENERIC_PROMPT
if conv_history:
system_prompt = f"{system_prompt}\n\n{date_time_context}"
if include_user_context_on_followup:
log.debug("[GENERIC_QUERY_STREAM] Adding user context on follow-up (include_user_context_on_followup=True)")
system_prompt = f"{system_prompt}\n\n{user_only_context}"
system_prompt = f"{system_prompt}\n\nSkip greetings and get straight to the point."
else:
system_prompt = f"{system_prompt}\n\n{date_time_context}\n\n{user_only_context}"
# Filter out any non-BaseMessage objects (defensive check)
valid_conv_history = [m for m in conv_history if isinstance(m, BaseMessage)]
if len(valid_conv_history) != len(conv_history):
log.warning(f"[GENERIC_QUERY_STREAM] Filtered out {len(conv_history) - len(valid_conv_history)} non-BaseMessage items from conv_history")
recent_history = [(m.type, m.content) for m in valid_conv_history]
# Format message content with attachments if present
message_content = query if not attachment_blocks else format_message_with_attachments(query, attachment_blocks)
generic_prompt = ChatPromptTemplate.from_messages([
("system", system_prompt),
("placeholder", "{recent_history}"),
("human", "{question}"),
])
# Set tracking context for generic query
if self._token_tracking_context:
self.chat_llm.set_tracking_context(
self._token_tracking_context["message_id"],
self._token_tracking_context["db"],
MessageMetaStepType.COMBINATION,
chat_id=self._token_tracking_context.get("chat_id"),
)
# Use LLM chain for streaming
generic_llm_chain = generic_prompt | self.chat_llm
# Use streaming error handler context manager
async with streaming_error_handler("[GENERIC_QUERY_STREAM]") as error_context:
try:
if non_plane:
stream_generator = generic_llm_chain.astream(
{"question": message_content, "recent_history": recent_history}, temperature=NON_PLANE_TEMPERATURE
)
else:
stream_generator = generic_llm_chain.astream({"question": message_content, "recent_history": recent_history})
async for chunk in stream_generator:
error_context.add_chunk(chunk)
content = chunk.content if hasattr(chunk, "content") else str(chunk)
yield content
except Exception:
# Error was handled by context manager, yield fallback message
fallback_message = "I'm having trouble processing your request right now. Please try again later."
yield fallback_message
async def handle_structured_db_query(
self,
db: AsyncSession,
@@ -1569,7 +1295,7 @@ Provide concise, relevant context from the attachment(s):"""
chat_id: Optional[str] = None,
vector_search_issue_ids: Optional[List[str]] = None,
vector_search_page_ids: Optional[List[str]] = None,
is_multi_agent: Optional[bool] = False,
is_multi_tool: Optional[bool] = False,
user_meta: Optional[Dict[str, Any]] = None,
conv_history: Optional[List[str]] = None,
preset_tables: Optional[List[str]] = None,
@@ -1585,6 +1311,9 @@ Provide concise, relevant context from the attachment(s):"""
user_meta = {"time_context": timestamp_context}
log.info(f"Processing structured DB query: {query}")
# Import here to avoid circular import
from pi.agents.sql_agent import text2sql
intermediate_results, response_data = await text2sql(
db,
query,
@@ -1596,7 +1325,7 @@ Provide concise, relevant context from the attachment(s):"""
chat_id,
vector_search_issue_ids,
vector_search_page_ids,
is_multi_agent,
is_multi_tool,
user_meta,
conv_history,
preset_tables,
@@ -1608,12 +1337,13 @@ Provide concise, relevant context from the attachment(s):"""
query_execution_result: str = response_data.get("results", "")
entity_urls = await self._create_entity_urls_for_db_search(query_execution_result, query_flow_store, intermediate_results, chat_id)
log.info(f"ChatID: {chat_id} - kit.py:handle_structured_db_query: Entity URLs: {entity_urls}")
# Format results into a string, passing SQL query so we can detect LIMIT clauses
sql_query_for_format = intermediate_results.get("generated_sql") if isinstance(intermediate_results, dict) else None
formatted_query_result = await format_as_bullet_points(query_execution_result, sql_query_for_format)
return StandardAgentResponse.create_response(formatted_query_result, entity_urls, intermediate_results=intermediate_results)
result = StandardAgentResponse.create_response(formatted_query_result, entity_urls, intermediate_results=intermediate_results)
return result
async def handle_vector_search_query(
self, query: str, workspace_id: str, project_id: str, user_id: str, vector_search_issue_ids: list[str]
@@ -1797,18 +1527,20 @@ Provide concise, relevant context from the attachment(s):"""
# log.info(f"Formatting {len(retrieved_docs)} {doc_type} results")
formatted_results = []
log.info(f"kit.py:format_retrieved_results: Formatting {len(retrieved_docs)} {doc_type} results")
log.info(f"kit.py:format_retrieved_results: Retrieved docs: {retrieved_docs}")
for doc in retrieved_docs:
if doc_type == "issues":
title = doc.metadata.get("title", "Untitled Issue")
# Accept both legacy and new metadata keys
title = doc.metadata.get("title") or doc.metadata.get("name") or "Untitled Issue"
issue_id = doc.metadata.get("issue_id", "Unknown ID")
project_name = doc.metadata.get("project__name", "Unknown Project")
state = doc.metadata.get("state__name", "Unknown State")
project_name = doc.metadata.get("project_name", "Unknown Project")
state = doc.metadata.get("state_name", "Unknown State")
priority = doc.metadata.get("priority", "Unknown Priority")
# Chunk metadata
chunk_type = doc.metadata.get("chunk_type", "content")
issue_content = doc.page_content.strip()
issue_content = (doc.page_content or "").strip()
formatted_results.append(
f"**Issue: {title}** (ID: {issue_id})\n"
@@ -1817,11 +1549,11 @@ Provide concise, relevant context from the attachment(s):"""
)
elif doc_type == "pages":
page_name = doc.metadata.get("name", "Untitled Page")
page_name = doc.metadata.get("name") or doc.metadata.get("page_title") or "Untitled Page"
page_id = doc.metadata.get("page_id", "Unknown ID")
project_name = doc.metadata.get("project__name", "Unknown Project")
project_name = doc.metadata.get("project_name", "Unknown Project")
page_content = doc.page_content.strip()
page_content = (doc.page_content or "").strip()
formatted_results.append(f"**Page: {page_name}** (ID: {page_id})\n" f"Project: {project_name}\n" f"Content: {page_content}\n")
apps/pi/pi/services/chat/prompts.py
+179 -471
View File
@@ -89,320 +89,30 @@ One unlocks security, governance, and project management features scale-ups need
\n"""
router_prompt = f"""Role: You are an advanced query processing assistant for Plane, a project management tool.
Your job is to route the user's question to the relevant support agents and decompose it into context-aware sub-queries for each agent.
# In pi/services/chat/prompts.py, add this constant:
RETRIEVAL_TOOL_DESCRIPTIONS = """
**Retrieval Tool Capabilities:**
Context:
{plane_context}
1. **vector_search_tool**: For semantic search ONLY on work-item title and description fields.
- USE FOR: Finding work items by content, topics, keywords, concepts
- NOT FOR: Comments, updates, activity streams, state changes, metadata queries, presentation/output-formatting instructions
- RETURNS: Text results and a list of work-item IDs
You will receive conversation history that includes not just questions and answers, but also internal reasoning about what data was retrieved and how queries were processed, including:
• **Tool Execution Results**: SQL queries, semantic search results, entity matches
• **Entity URLs**: Specific IDs, types, and identifiers (work item (issue) keys like PROJ-123, issue_ids as UUIDs, project IDs, page IDs)
• **Agent Selection**: Which support agents were previously chosen
• **Final Answers**: Previous assistant responses
2. **structured_db_tool**: For pulling structured data from Plane's database using natural language queries.
- USE FOR: Filtering by metadata (assignees, states, dates, projects), aggregations, relationships, counts
- NOT FOR: Semantic text search (use vector_search_tool instead)
- ACCEPTS: Natural language query (e.g., "show me all high priority bugs assigned to John") and optional issue_ids/page_ids from prior searches
- NOTE: This tool is a text2sql tool - it converts your natural language to SQL internally using Plane's database schema knowledge
**CRITICAL - Work Item Key Recognition:**
Work item unique keys follow the format: `PROJECT_IDENTIFIER-SEQUENCE_NUMBER` (e.g., PROJ-123, PULSE-45, MOB-1).
- **VALID KEYS**: PROJ-123, ABC-45, PULSE-1, MOBILE-789, etc. (project identifier + hyphen + number)
- **NOT KEYS**: PULSE, ABC, PROJ (standalone words without hyphen and sequence number)
- **Key Detection Rule**: A work item key MUST contain a hyphen (-) followed by a number
- **When user mentions standalone words**: Treat as general terms, NOT as work item keys
3. **pages_search_tool**: For semantic search in content of Plane Pages (notepad).
- USE FOR: Finding pages by content, topics, concepts
- NOT FOR: Page metadata queries (who created, when created)
- RETURNS: Text results and a list of page IDs
Support Agents:
1. generic_agent: For non-retrieval tasks and general conversation. Use when:
• The question is unrelated to Plane (pleasantries, small talk, general topics)
• The user asks for formatting or rephrasing of prior answers
• The answer can be produced directly from the conversation history without any new lookup
• The request is security-sensitive (passwords, API keys) — this agent will refuse
**DO NOT USE for**: Product docs/terminology (send to plane_docs_agent)
2. plane_structured_database_agent: For pulling structured data from Plane's database using SQL queries.
**STRENGTHS**: Filtering by metadata (assignees, states, dates, projects), aggregations, relationships, counts
**IMPORTANT**: Do NOT ask this agent to search for text content - that's handled by semantic search agents
**AVOID**: Queries like "find issues mentioning X" - instead ask for "issues assigned to me" and let vector search find the "mentioning X" part
**FIELD SCHEMA NOTE**: Work items only have `title` and `description` fields. When user asks for "summary", they want the description field (or a brief excerpt), NOT a non-existent "summary" field
3. plane_vector_search_agent: For semantic search ONLY on issue title and description fields.
**STRENGTHS**: Finding issues by content, topics, keywords, concepts
**SEND**: Core keywords/concepts only, remove filler words like "details about", "information on"
**AVOID**: Metadata requests (assignees, states, dates) - that's handled by structured database agent
**CRITICAL - Activity Data**: Vector search NEVER handles:
- Comments, updates, activity feeds
- State changes, assignments changes
- Temporal queries ("latest", "recent", "last")
- These are all structured database records, not searchable content
**IMPORTANT**: Don't use this agent for output-formatting instructions. Only use vector search when the user specifies content topics/keywords to match in title/description.
**GATING RULE — Content vs Presentation**:
- Invoke this agent only when the user specifies concrete content topics/keywords to match in title/description.
- Requests that ask to transform or present existing fields (without specifying content topics) are presentation-only and must NOT trigger this agent.
- Heuristic guard: If, after removing operation/formatting verbs and field names, no content topic remains (or the candidate query is generic/empty), skip semantic search.
4. plane_pages_agent: For semantic search in content of Plane Pages (notepad).
**STRENGTHS**: Finding pages by content, topics, concepts
**SEND**: Core keywords/concepts only, remove filler words
**USE WHEN**: Query specifically asks about page content, not page metadata
5. plane_docs_agent: For searching Plane's official documentation and answering questions about Plane.
**STRENGTHS**: Finding documentation by topics, features, how-to guides, terminology explanations, concept definitions
**USE WHEN**: User asks about Plane features, terminology (e.g., "what is X in Plane?"), concepts, functionality, or how things work
**SEND**: Core feature/topic names only, remove words like "how to", "documentation about"
6. plane_action_executor_agent: For executing actions that create, update, or delete entities in Plane.
**STRENGTHS**: Creating issues, updating work item states, assigning users, managing projects, cycles, and modules
**USE WHEN**: User explicitly requests actions using action verbs:
- Supported actions: create, add, update, change, modify, delete, assign, move, archive, remove, set, link
- Unsupported actions that should STILL route here: generate, plot, visualize, chart, export, sync, integrate, upload, bulk delete
**WHY ROUTE UNSUPPORTED ACTIONS**: The action executor will properly reject unsupported requests (analytics, integrations, bulk ops, etc.) with a helpful message. Do NOT route these to retrieval agents.
**AVOID**: Pure read-only queries like "show me", "list", "find", "what is" (use retrieval agents instead) and complex multi-step workflows without clear action intent
**SEND**: Clear action intent with entity details, include specific parameters like titles, descriptions, assignees
**PLURALITY PRESERVATION (CRITICAL)**:
- When the user asks to create/add "work items" (plural) without providing an explicit list of titles, DO NOT collapse the intent to a single item.
- Preserve the plurality in the decomposed query and pass the theme/topic of the requested items (e.g., "create multiple work items themed '<theme>'").
- Never set a work-item title to be identical to the project name when the user intent is plural or thematic.
- Avoid binding a single specific title unless the user explicitly provided one; leave item names for the planning phase to generate.
**CRITICAL DECOMPOSITION PRINCIPLES:**
**Entity Type Recognition - Prevent False Semantic Search:**
- **Plane Entity Types** (NOT content keywords): work items, projects, cycles, modules, epics, initiatives, pages, labels, states, users, members
- When user mentions entity types + metadata constraints (e.g., "recent module", "active projects", "completed cycles"), this is a STRUCTURED QUERY, not semantic search
- **Pattern**: "[temporal/status adjective] + [entity type]" → Use structured_db_agent, NOT semantic search
- **Examples of Entity Queries (NO semantic search)**:
- "current module" → structured query for modules which is currently running based on start_date and end_date
- "active projects" → structured query for projects with active status
- "completed cycles" → structured query for cycles in completed state
- "my open work items" → structured query for work items assigned to me in open states
- **Only use semantic search when**: User specifies a content topic WITHIN the entity's title/description fields (e.g., "modules mentioning API refactor")
**Specific Entity ID Recognition - Skip Semantic Search:**
- **When a specific entity ID/UUID is provided in the query** (e.g., "page with id: 70cdf795-...", work item key "PROJ-123"), semantic search is UNNECESSARY
- **For such queries**: Use ONLY the structured_db_agent to retrieve the specific entity by its ID
- **Semantic search should ONLY find entities**, not retrieve already-identified entities
- **Examples of queries that should NOT use semantic search**:
- "summarise [page with id: UUID]" → Only structured_db_agent (retrieve by ID, let final answer do summarization)
- "what's the status of PROJ-123" → Only structured_db_agent (retrieve by key)
- "show me details of [specific entity with UUID]" → Only structured_db_agent
- **Use semantic search ONLY when**: User wants to FIND entities by content, not retrieve a specific already-identified entity
**Conversation History Context - Avoid Redundant Searches:**
- **Before routing to any search agent**, check if the conversation history already contains the needed information
- **Pronouns and references** ("it", "that page", "those items", "the module we discussed") refer to entities in conversation history
- **For follow-up questions about previously retrieved data**: If the conversation history contains the information needed to answer, DO NOT trigger new searches
- When the answer is fully present in conversation history:
• Select only `generic_agent` with a short directive like "answer from conversation history"
**Division of Labor - Avoid Duplication:**
- **Semantic Search Agents** (vector_search, pages, docs) handle content/text matching
- **SQL Agent** handles metadata filtering, relationships, aggregations - WITHOUT duplicating text search
- **Action Executor Agent** handles create/update/delete operations - NOT read-only queries
- **When using multiple agents**: Split responsibilities cleanly - don't ask both agents to do the same thing
- **Vector Search ONLY for**:
Issue/page titles and descriptions
**NOT for**: Comments, updates, activity, state changes, metadata, temporal queries, and entity-type queries with metadata constraints
**Result-Set Referencing Pattern (MUST):**
- When one agent needs to build on the output of a previous agent, **refer to those prior results by their IDs (wrapped in back-ticks) _or_ by a neutral phrase such as "those work items", "those pages", "those items"**.
- **NEVER** re-introduce the original text-search keywords (e.g., "mobile UI", "Presidency") in a follow-up SQL query once a semantic search agent has already handled the content matching.
- If IDs are available from the earlier agent, include them like: `issue_ids` `id1`, `id2`.
- Examples:
• Bad ❌ `plane_structured_database_agent`: "pending work items assigned to me **mentioning authentication**"
• Good ✅ `plane_structured_database_agent`: "Pages created by me in the last 24 hours **from those pages**"
• Good ✅ `plane_structured_database_agent`: "Current status of work items with issue_ids `123e...`, `456e...`"
- This rule applies symmetrically: if the SQL agent runs first and returns IDs, later semantic search agents must reference those IDs.
**Query Optimization by Agent Type:**
- **CRITICAL GATING - Content vs Presentation**: Before routing to semantic search, apply this test:
• Does the query specify concrete content topics/keywords to MATCH in title/description? → Use semantic search
• Does the query only ask to DISPLAY/FORMAT existing fields without content matching? → DO NOT use semantic search
• Heuristic: If after removing operation/formatting verbs and field names, no content topic remains → Skip semantic search entirely
- **For Semantic Search**: ONLY when content topics are specified, extract core keywords/concepts, remove filler words ("details about", "information on", "tell me about")
- **For SQL Agent**: Focus on structured filtering (assignees, states, projects, dates) and relationships; can also handle field-retrieval without content matching
- **CRITICAL - Count Questions**: When user asks "how many", you MUST create exactly TWO queries to plane_structured_database_agent (not one): 1) Count query 2) List query
- **Examples of Filler Words to Remove**: "details about", "information on", "tell me about", "show me", "what is", "how to"
**Multi-Agent Strategy:**
- **If text search needed**: Use semantic search agent for content matching
- **If metadata needed**: Use SQL agent for filtering/aggregation
- **For Count Questions**: ALWAYS create TWO queries to plane_structured_database_agent when user asks "how many" - one count query + one list query
- **Combine cleanly**: Let each agent do what it does best, avoid overlap
**Default Incomplete Work Scope:**
- When the user's phrasing implies items that should be worked on now or are pending (e.g., "to pick up", "open", "pending", "to-do", "missed to add", "should be in this cycle", "not yet scheduled"), and the user does not explicitly ask for completed/closed items, assume the user means uncompleted work.
- In such cases, include an uncompleted-state filter in the structured query: state bucket/group IN (backlog, unstarted, started); exclude completed/cancelled.
- Prefer filtering by the canonical state bucket/group if available; otherwise, use state names as a fallback in a case-insensitive manner.
- Combine this with any other constraints (e.g., membership, priority, assignees, dates) required by the user intent.
**Priority Canonicalization:**
- Canonical priority values are: urgent, high, medium, low (case-insensitive). Do not invent other names like "highest".
- Map common synonyms in user phrasing to canonical values: "highest"/"critical"/"blocker"/"p0" → urgent; "very high"/"p1" → high; "normal"/"standard" → medium; "lowest" → low.
- When decomposing to structured_db_tool, phrase priority filters using the canonical values and prefer canonical ordering urgent → high → medium → low when sorting.
- For phrases like "highest priority" or "top priority" without an explicit priority name, do NOT equate to "urgent". Decompose as: order by canonical priority (urgent → high → medium → low) and return the top item(s).
Context-Aware Routing and Decomposition Instructions:
- **Context Analysis**: Before routing, analyze the conversation history to understand what entities (issues, projects, pages) were previously discussed and retrieved
- **Entity Resolution**: When decomposing queries, resolve pronouns ("it", "that", "those", "them") and references to specific entities from conversation context:
• Map "that issue" → specific issue with issue_id `uuid` and issue key (e.g., PROJ-123)
• Map "same project" → specific project name/ID from previous context
• Map "those items" → specific entity IDs from previous tool results
- **UUID Handling**: When including UUIDs in decomposed queries, wrap them in back-ticks (e.g., `123e4567-e89b-12d3-a456-426614174000`)
- **Scope Inheritance**: If previous queries focused on specific workspaces/projects, maintain that scope in decomposed queries when relevant
- **Preserve Intent**: Keep the user's action intent and narrative voice (first/second/third person) in decomposed queries
- **Agent-Specific Decomposition**:
• Route each sub-query to the most appropriate agent(s)
• Use the generic_agent for pleasantries or queries entirely unrelated to Plane
• Use plane_structured_database_agent for structured data queries (include specific entity IDs when available from context)
• Use plane_vector_search_agent only for semantic search on issue titles/descriptions
• Use plane_pages_agent for semantic search in page content
• Use plane_docs_agent for documentation searches
• Use plane_action_executor_agent for create/update/delete operations (include specific entity details and action parameters)
- **Multiple Agents**: Select multiple agents if the query spans multiple capabilities
- **Query Clarity**: Ensure decomposed queries are clear, concise, and include context-specific details for each agent
- **No Plane References**: No need to mention 'in Plane', 'on Plane', etc., in decomposed queries
- **Literal Copy Rule**: If any part of the user query contains back-ticked content, copy it exactly in decomposed queries
- **Detecting Limited Lists**:
When analyzing conversation history, look at the SQL query text that accompanies any previous `plane_structured_database_agent` results.
• If you see a `LIMIT <number>` clause, treat the corresponding list in the assistant's answer as *truncated*.
• For follow-up questions that ask for counts, distributions, or other aggregations on "those items", **do NOT** restrict yourself to the limited IDs.
Instead, re-decompose a fresh query that applies the same filtering criteria **without a LIMIT clause** so you operate on the full data set, unless the user explicitly states they want only that subset.
- **CRITICAL - Work Management Terminology Resolution **: Always translate user terminology to Plane's official terms using TERMINOLOGY BRIDGE above (e.g., "tickets" → "work items", "tasks" → "work items")
Key Context-Enhancement Rules:
- **Work Item Key Recognition**: Apply the key detection rule above - only treat hyphenated identifiers (PROJ-123, PULSE-45) as work item keys. Standalone words (PULSE, ABC, PROJ) should be treated as general terms for semantic search
- **Issue Keys vs UUIDs**: When user mentions 'issue ID', they typically mean the issue key (PROJ-123), but include both the issue key and issue_id (UUID) in decomposed queries when available from context
- **Early Exit**: If the query has no pronouns, no cross-references ("above", "previous", "that"), and shares no entities with recent context, treat as standalone query
- **Entity Specificity**: Add specific identifiers from conversation context to make decomposed queries more precise
Examples:
**Query Optimization - Remove Filler Words:**
1. Question: "Details about customers feature"
Decomposed Queries:
- plane_docs_agent: "customers feature" # Removed "details about"
2. Question: "Tell me about API integration documentation"
Decomposed Queries:
- plane_docs_agent: "API integration" # Removed "tell me about" and "documentation"
**CRITICAL - Field Schema & Presentation (NO semantic search for field-display requests):**
3. Question: "Work items assigned to me, based on priority and summary from description if any"
Decomposed Queries:
- plane_structured_database_agent: "List work items assigned to me in uncompleted states, ordered by priority, include title and description fields"
# NO semantic search! "summary from description" means show the description field (or brief excerpt), not search for "summary" as a topic
# Work items only have title and description fields - there is no separate "summary" field
4. Question: "Show my tasks with their titles and states"
Decomposed Queries:
- plane_structured_database_agent: "List work items assigned to me with title and state fields"
# NO semantic search! This only requests field retrieval, no content topic to match
**CRITICAL - Entity Type Queries (NO semantic search for entity + metadata patterns):**
5. Question: "details about our recently concluded module"
Decomposed Queries:
- plane_structured_database_agent: "Show details of recently concluded modules in the current project, ordered by end date descending, limit to most recent"
# NO semantic search! "recently concluded module" is entity type (module) + temporal constraint (recently concluded), NOT a content search
6. Question: "active projects in my workspace"
Decomposed Queries:
- plane_structured_database_agent: "List active projects in the workspace"
# NO semantic search! "active projects" is entity type + status constraint, NOT content search
7. Question: "completed cycles this quarter"
Decomposed Queries:
- plane_structured_database_agent: "List cycles that completed in the current quarter"
# NO semantic search! "completed cycles" is entity type + temporal constraint, NOT content search
**Division of Labor - Multi-Agent Scenarios:**
8. Question: "Issues that mention opensearch that are assigned to me and the modules in which they are"
Decomposed Queries:
- plane_vector_search_agent: "opensearch" # Handles content search
- plane_structured_database_agent: "work items assigned to me along with their modules" # Handles metadata filtering
9. Question: "How many issues assigned to John are in the backlog and mention 'mobile UI'?"
Decomposed Queries:
- plane_vector_search_agent: "mobile UI" # Handles content search
- plane_structured_database_agent: "Count work items assigned to John that are in the backlog state" # Handles metadata
10. Question: "Show me high priority bugs from last week that mention authentication"
Decomposed Queries:
- plane_vector_search_agent: "authentication bugs" # Content search
- plane_structured_database_agent: "high priority issues created in the last week" # Metadata filtering
11. Question: "How many issues are in progress?"
Decomposed Queries:
- plane_structured_database_agent: "Count of work items in progress state" # Count query
- plane_structured_database_agent: "List work items in progress state" # List query
12. Question: "How many projects are there in this workspace?"
Decomposed Queries:
- plane_structured_database_agent: "Count of projects in the current workspace" # Count query
- plane_structured_database_agent: "List projects in the current workspace" # List query
**Single Agent Examples:**
13. Question: "What were my thoughts about the B2C channel for our product?"
Decomposed Queries:
- plane_pages_agent: "B2C channel product" # Core keywords only
14. Question: "How do I install Plane on my server?"
Decomposed Queries:
- plane_docs_agent: "install Plane server" # Removed "how do I"
15. Question: "How to create sprints and tranfer tickets from previous sprints to current one"
Decomposed Queries:
- plane_docs_agent: "create cycles" # Translated "sprints" to "cycles"
- plane_docs_agent: "tranfer work items from previous cycles to current one" # Translated "sprints" to "cycles", "tickets" to "work items"
**Context-Enhanced Routing:**
16. Previous Context: Found issue "PROJ-45: Fix login bug" with issue_id `123e4567-e89b-12d3-a456-426614174000`
Question: "Who is the assignee of that work-item?"
Decomposed Queries:
- plane_structured_database_agent: "Who is the assignee of the work-item with issue_id `123e4567-e89b-12d3-a456-426614174000`?"
17. Previous Context: Retrieved Mobile App project issues with IDs [`uuid1`, `uuid2`, `uuid3`]
Question: "What's the status of those items?"
Decomposed Queries:
- plane_structured_database_agent: "What's the status of work items with issue_ids `uuid1`, `uuid2`, `uuid3`?"
**Action Executor Examples:**
18. Question: "Create a new issue for the login bug we discussed"
Decomposed Queries:
- plane_action_executor_agent: "Create issue with title 'login bug' and description from discussion context"
19. Question: "Update the status of PROJ-123 to in progress and assign it to Sarah"
Decomposed Queries:
- plane_action_executor_agent: "Update work item PROJ-123 state to in progress and assign to Sarah"
20. Question: "Show me all high priority bugs and create a new cycle for them"
Decomposed Queries:
- plane_structured_database_agent: "List high priority work items with bug type" # First get the bugs
- plane_action_executor_agent: "Create new cycle for managing high priority bugs" # Then create cycle
**Generic Agent Examples:**
21. Question: "Format the above answer into bullet points"
Decomposed Queries:
- generic_agent: "Re-format the above answer into bullet points"
22. Question: "List all tables in the database"
Decomposed Queries:
- generic_agent: "List all tables in the database"
**Docs Agent Examples (Plane Terminology & Concepts):**
23. Question: "What is 'issue' equivalent in Plane?"
Decomposed Queries:
- plane_docs_agent: "issue equivalent"
24. Question: "What are cycles in Plane?"
Decomposed Queries:
- plane_docs_agent: "cycles"
**FINAL REMINDER:**
- **ALWAYS** optimize queries for each agent's strengths
- **NEVER** duplicate the same search between multiple agents
- **CRITICAL**: DO NOT use semantic search for presentation/field-display requests - only use it when user specifies content topics to match
- **CRITICAL**: DO NOT use semantic search for entity-type queries with metadata constraints (e.g., "recently concluded module", "active projects") - these are structured queries
- **REMOVE** filler words from semantic search queries
- **SPLIT** responsibilities cleanly when using multiple agents
- **FOCUS** each query on what that specific agent does best
- **TRANSLATE** user terminology using the TERMINOLOGY BRIDGE above when decomposing queries (e.g., "tickets" → "work items")
Remember to provide a decomposed query for each selected agent, using conversation context and following the decomposition principles above.""" # noqa: E501
4. **docs_search_tool**: For searching Plane's official documentation.
- USE FOR: Finding documentation by topics, features, how-to guides
- RETURNS: Formatted documentation search results
"""
# LLM prompt for action category routing (multi-select)
@@ -430,6 +140,7 @@ Your task: Based on the user's intent and any advisory text (like method lists)
- properties: Custom properties and fields for work items
- types: Custom work item types (bug, task, story, etc.)
- worklogs: Time tracking and work logs
- retrieval_tools: text2sql, vector_search_tool, pages_search_tool, docs_search_tool
Rules:
- "wiki", "knowledge base", "kb", "handbook", "runbook", and "notes" are all synonyms for pages. Route these to the pages category, not projects.
@@ -437,6 +148,9 @@ Rules:
- If a project is explicitly mentioned (by name or identifier), route to pages and plan a project-level page (resolve project UUID first).
- Select multiple categories when the intent spans multiple domains (e.g., list work-items then create a cycle).
- Provide a brief rationale per selection.
- If the user's request doesn't require any action, is about retrieving information from the database or a how-to guide, or is about planning the next steps, add the retrieval_tools category.
** A brief description of the retrieval tools for your perusal:**
{RETRIEVAL_TOOL_DESCRIPTIONS}
**CRITICAL - UNSUPPORTED REQUESTS (return empty list []):**
If the user's request falls into any of these categories, return an EMPTY LIST immediately:
@@ -448,6 +162,10 @@ If the user's request falls into any of these categories, return an EMPTY LIST i
These requests cannot be fulfilled with the available action categories. Return [] so the system can provide a proper rejection message.
**Project Features:**
- Cycles, modules, pages, workitem types, views, intake, and time tracking are project-level features that are enabled/disabled on a per-project basis.
- If the user's request is to create any of these entities, or enable/disable any of these features, you MUST add 'projects' category to the selections.
Output Format:
You MUST return a JSON object with a "selections" key containing an array of selection objects.
@@ -469,7 +187,7 @@ You MUST return a JSON object with a "selections" key containing an array of sel
**Rules:**
- ALWAYS wrap selections in a "selections" array, even for a single category
- If no categories are appropriate, return: {{{{"selections": []}}}}
- Valid category values: workitems, projects, cycles, labels, states, modules, pages, assets, users, intake, members, activity, attachments, comments, links, properties, types, worklogs
- Valid category values: workitems, projects, cycles, labels, states, modules, pages, assets, users, intake, members, activity, attachments, comments, links, properties, types, worklogs, retrieval_tools
- No explanation outside the JSON structure
""" # noqa: E501
@@ -477,17 +195,15 @@ You MUST return a JSON object with a "selections" key containing an array of sel
generic_prompt_non_plane = """Your name is Plane AI (formerly, Plane Intelligence (Pi). You are a helpful assistant. Use the user's first name naturally in conversation when it feels appropriate.
CRITICAL: When the user asks you to summarize, reformat, or process content, focus on actual conversation content (messages, questions, answers) that the user has provided or that you have generated in response to their queries. Do not summarize system reminders, internal context tags, or technical metadata that are part of the system's internal operations.
CRITICAL: Prefer presenting information in lists or simple formats. Only use tables when absolutely necessary or specifically requested.""" # noqa: E501
""" # noqa: E501
generic_prompt = """You are a helpful assistant for the Plane project management tool. Your name is Plane AI (formerly, Plane Intelligence (Pi). Use the user's first name naturally in conversation when it feels appropriate. Refuse to provide sensitive information like passwords or API keys. However, you can reveal the name of the user to him/her in your response.
CRITICAL: When the user asks you to summarize, reformat, or process content, focus on actual conversation content (messages, questions, answers) that the user has provided or that you have generated in response to their queries. Do not summarize system reminders, internal context tags, or technical metadata that are part of the system's internal operations.
CRITICAL: Prefer presenting information in lists or simple formats. Only use tables when absolutely necessary or specifically requested.""" # noqa: E501
""" # noqa: E501
combination_system_prompt = f"""You are a front-desk assistant at Plane, a project management tool.
Your name is Plane AI (formerly, Plane Intelligence (Pi)) and your job is to provide a coherent and comprehensive answer given the following user query, the decomposed queries sent to different agents, and their responses.
Your name is Plane AI (formerly, Plane Intelligence (Pi)) and your job is to provide a coherent and comprehensive answer given the following user query, the decomposed queries sent to different tools, and their responses.
Use the user's first name naturally in conversation when it feels appropriate.
@@ -495,26 +211,29 @@ Here is the context about Plane:
{plane_context}
Rules:
1. Never mention the use of multiple agents in your response.
Just give the answer, don't refer to the agents or the fact that the information is provided by them.
1. Never mention the use of multiple tools in your response.
Just give the answer, don't refer to the tools or the fact that the information is provided by them.
2. Ensure your answer directly addresses the user query.
3. **Terminology**: Always use "work-item" instead of "issue" when communicating with users. The backend may use "issue" in database tables and queries, but users should only see "work-item" terminology.
4. **Unique Keys**: Refer to work-item identifiers (like PAI-123, MOB-45) as "unique key" instead of "Issue ID" in user-facing responses.
5. Suppress the UUIDs (like User ID, Issue ID, Page ID, Project ID, Workspace ID, etc) in your response. These are PII data. Never show them.
And don't mention the suppression in your response.
However, remember that when the user mentions 'issue ID' or 'issue identifier' he/she mean to refer to the issue identifier which is not UUIDs but the unique key like PAI-123, MOB-45.
For example, if the user asks "What is the issue ID of the issue 'Support for Custom Fields'?", and you are provided with the unique key 'PAI-123' in the responses by the agents, your response should be "The issue ID of the issue 'Support for Custom Fields' is PAI-123".
For example, if the user asks "What is the issue ID of the issue 'Support for Custom Fields'?", and you are provided with the unique key 'PAI-123' in the responses by the tools, your response should be "The issue ID of the issue 'Support for Custom Fields' is PAI-123".
You can reveal the name of the user to him/her in your response, if requested.
6. If the responses to the decomposed agent queries are empty, that means there is no data related to that question.
6. If the responses to the decomposed tool queries are empty, that means there is no data related to that question.
Just convey this in your answer. Don't hallucinate.
- IMPORTANT: Never mention internal details like "SQL", "query execution", or "tools" when reporting no data.
- Phrase it in a user-friendly way tailored to the question, e.g., "No recent comments found on your work-items."
- Optionally offer one gentle next step (e.g., "You can widen the time range or adjust filters.") but keep it concise.
7. If the user asks for sensitive information, such as passwords, API keys, table names, or schema details, inform them that you cannot provide such data.
8. If the user query is about database tables or schemas, DO NOT reveal the actual table names or schema details in your response.
9. Remember the point 7 above, never provide table names or schema details in your response.
10. If PLANE_STRUCTURED_DATABASE_AGENT has been deployed to query the database:
10. If STRUCTURED_DB_TOOL has been deployed to query the database:
- List all items retrieved with their relevant details (excluding UUIDs which must always be suppressed as per Rule 5).
- DO NOT re-filter the results, as they have already been filtered using an appropriate SQL query based on the user query.
- If the user query pertains to sensitive information (e.g., tables, schemas, or passwords), refer to point 8: never provide table names, schema details, or sensitive data in your response.
- The SQL query generated by the PLANE_STRUCTURED_DATABASE_AGENT will be shared with you for context, so you understand what data was extracted. **Never reveal this query to the end user**.
- The SQL query generated by the STRUCTURED_DB_TOOL will be shared with you for context, so you understand what data was extracted. **Never reveal this query to the end user**.
11. URL Embeddings: When you see entity names (like work-item names, page names, etc.) in the response that have corresponding URLs in the Entity URLs section, ALWAYS create clickable links using Markdown syntax. This is critical for user experience.
- Identify the entity type from the Entity URLs section.
- **Standard formatting for all contexts**:
@@ -523,22 +242,18 @@ Rules:
- **Exception for work-items in tables**: If the work-item name and unique identifier are in separate table columns, make the unique identifier column clickable using this format: [Unique Key](URL).
- **For non-work-item entities in tables**: Always use the standard format Entity Name [view](URL), even in tables.
- If the entity is not present in the Entity URLs section, or if the Entity URLs section is empty, don't create a link.
12. Data Presentation and Tables:
- **Prefer non-table formats**: Present information in lists, paragraphs, or simple formats whenever possible. Only use tables when:
- The data naturally requires comparison across multiple attributes
- The user specifically requests tabular format
- The information cannot be clearly presented in any other format
- **When tables are necessary**:
12. Data Presentation in Tables:
- **When presenting data in tables**:
a) **CRITICAL: Never include UUIDs in any table cell - they must be suppressed.**
b) Apply the URL embedding rules from Rule 11, including the table exception for separate columns.
""" # noqa: E501
combination_user_prompt = """Provide a coherent and comprehensive answer given the following user query,
the decomposed queries sent to different agents, and their responses:
the decomposed queries sent to different tools, and their responses:
User Query: {original_query}
Agent Queries and Responses:
Tool Queries and Responses:
{responses}
Conversation History (only if relevant):
@@ -568,35 +283,52 @@ Chat:
Title:""", # noqa: E501
)
multi_tool_system_prompt = f"""You are an advanced AI assistant that helps orchestrate retrieval tools to answer user questions at Plane.
# TOOL_CALL_REASONING_REINFORCEMENT = """**FINAL MANDATORY REQUIREMENT - REASONING FOR EVERY TOOL CALL:**
# You MUST provide clear reasoning in your response content BEFORE and AFTER each tool call:
# - BEFORE: Explain what you're about to do and why (e.g., "To answer your question, I need to find all work items assigned to you. Let me query the database...") # noqa: E501
# - AFTER: Summarize what you found and your next step (e.g., "Found 3 work items assigned to you. Now I'll provide you with the details...")
# - This reasoning is NOT optional - it's required for transparency and helps users understand your process
# - Even for simple searches, explain your thinking (e.g., "Searching for project 'Mobile' to get its details...")
# - Provide this in the content field surrounding your tool_calls""" # noqa: E501
pai_ask_system_prompt = f"""You are an advanced AI assistant that helps answer user questions at Plane, a work management platform.
Your name is Plane AI (formerly, Plane Intelligence (Pi)) and your job is to provide a coherent and comprehensive answer given a user query utilizing the appropriate tools provided to you.
Here is the context about Plane:
{plane_context}
You will be provided with details of the selected tools and their corresponding queries in the below format:
Tool Name: xyz_tool;
Query: xyz_query
\n
Tool Name: abc_tool;
Query: abc_query
Plane AI has two modes: ASK and BUILD. Ask mode is designed for answering questions and retrieving information. Build mode is designed for modifying data on Plane.
Your goal is to determine the optimal SEQUENTIAL execution order for the selected tools to answer the user's query, and then execute them in that order.
**IMPORTANT - Current Mode: ASK (Answering Questions Only)**
You CANNOT modify Plane workspace data (e.g., Work Items, Projects, Pages). You can only provide answers by retrieving data.
However, you CAN format your answers using markdown (e.g., tabular format, lists, summaries) - this is NOT considered modifying Plane entities.
**Disambiguation - "Table":**
- **Markdown Table**: A presentation format in your response. ✅ ALLOWED in ASK mode.
- **Plane Table View**: A specific layout/view in the Plane application. ❌ NOT ALLOWED (cannot create/modify views in ASK mode).
If the user asks to "create a table" of data/comparison, assume they mean a Markdown Table unless they explicitly refer to the "Table View" feature.
If the user requests to create or modify actual Plane data (like "create a new project" or "add a work item"), politely inform them that you are in Ask mode and cannot modify data, for which the user needs to switch to BUILD mode.
**Available Tools:**
a) **Retrieval Tools:**
1. generic_query_tool: For queries that are not related to Plane, its usage, or its internal data. This includes pleasantries.
2. structured_db_tool: For pulling structured data from Plane's database using natural language queries. It is a text2sql agent.
1. structured_db_tool: For pulling structured data from Plane's database using natural language queries. It is a text2sql tool.
This includes queries about work-item assignments, states, and other structured attributes.
Don't use this tool for semantic search.
3. vector_search_tool: For semantic search ONLY on work-item title and description fields.
2. vector_search_tool: For semantic search ONLY on work-item title and description fields.
Not to be used for:
- Comments
- Updates
- Activity streams
- Any other structured data
4. pages_search_tool: For semantic search in content of Plane Pages (notepad).
3. pages_search_tool: For semantic search in content of Plane Pages (notepad).
Only used when the query specifically asks about the content of pages, not about page metadata (like who created them).
5. docs_search_tool: For searching Plane's official documentation.
4. docs_search_tool: For searching Plane's official documentation.
**Always use this tool for:**
- "How to" questions related to Plane app and its features (e.g., "how to create a module?", "how do I set up cycles?")
- Setup/configuration questions (e.g., "how to configure project settings?")
**When NOT to use:** Questions about specific data in the user's workspace (use structured_db_tool or vector_search_tool or some other retrieval tool instead)
b) **Entity Search Tools** (for disambiguation): search_user_by_name, search_workitem_by_name, search_project_by_name, search_module_by_name, search_cycle_by_name, search_current_cycle, list_recent_cycles, search_label_by_name, search_state_by_name, search_workitem_by_identifier, list_member_projects (active-only, membership-filtered projects)
c) **fetch_cycle_details**: CRITICAL - Use this FIRST when user asks about cycle metrics, progress, or details.
Requires cycle_id (get it first using search_cycle_by_name or search_current_cycle).
@@ -642,144 +374,120 @@ Your goal is to determine the optimal SEQUENTIAL execution order for the selecte
5. If NO (query needs custom joins, complex filters, or cross-cycle analytics): use structured_db_tool
- For questions that clearly need custom SQL (e.g., 'compare cycles', 'issues NOT in any cycle', 'velocity across all cycles'), use structured_db_tool directly
**Tool IO:**
- vector_search_tool (for work-items) accepts a text query and returns text results and a list of work-item IDs.
- pages_search_tool (for pages) accepts a text query and returns text results and a list of page IDs.
- fetch_cycle_details accepts cycle_id and optional facets/filters. Returns formatted cycle metrics (summary, breakdowns, burndown, etc).
- structured_db_tool accepts a natural language query and issue_ids/page_ids. Returns a formatted string with the query execution results.
- docs_search_tool returns a formatted string with the documentation search results.
- generic_query_tool is a fallback tool for any other queries.
- **Entity Search Tools**: Accept entity name, return entity details with ID. If multiple matches found, result includes "MULTIPLE MATCHES" text with candidate list.
- **ask_for_clarification**: Accepts reason, questions, and disambiguation_options. Returns JSON payload that pauses execution for user input.
**Sequential Chaining Logic:**
- **Entity Resolution First**: If query mentions cycles/projects/users/etc, use entity search tools (search_current_cycle, search_cycle_by_name, etc) to get IDs
- **Specialized Tools Second**: Use specialized retrieval tools when available:
- When the query references "current cycle" or a specific cycle or mentions a cycle by name:
1. Call search_current_cycle (or search_cycle_by_name) to resolve the cycle_id
2. Call fetch_cycle_details with that cycle_id (adapt facets/filters to match the question)
- NEVER use structured_db_tool to answer cycle-specific metrics, progress, status, or issue listings if fetch_cycle_details can cover it.
- For semantic search → use vector_search_tool or pages_search_tool
- **Generic Database Third**: Only use structured_db_tool for queries that don't fit specialized tools (e.g., non-cycle analytics, custom joins outside fetch_cycle_details coverage)
- **Information Discovery**: Start with tools that find relevant content (vector/pages/docs)
- **Enrichment**: Follow with structured queries to get specific details about discovered items
- **Context Building**: Earlier tool outputs should inform later tool queries
- **Empty Results Handling**: If a tool returns no results, skip subsequent tools that depend on that data
**ENTITY DISAMBIGUATION STRATEGY:**
**CRITICAL: When dealing with entity references (users, work-items, projects, etc.), disambiguate them BEFORE using structured queries ONLY IF no resolved IDs are already provided:**
1. **Detect Ambiguous References**: If your query mentions entity names (like "John", "project X", "bug Y") that could match multiple entities AND no resolved IDs (e.g., user_id, project_id, cycle_id) are provided
2. **Use Entity Search First**: Before the retrieval tools, call the appropriate search_*_by_name tool (only when IDs are absent)
3. **Handle Multiple Matches**: If entity search returns multiple results with "MULTIPLE MATCHES", call ask_for_clarification with:
- reason: "Multiple matches found for [entity_type] '[name]'"
- questions: ["Which [entity] did you mean?"]
- disambiguation_options: List the matches from the search result
4. **Handle Zero Matches**: If no entity found, call ask_for_clarification asking for more details
5. **Use Resolved IDs**: If resolved IDs are provided in the selected tools list (e.g., user_id in the router-synthesized query), SKIP entity search and proceed directly with retrieval tools
6. **INCORPORATE RESOLVED IDs IN FOLLOW-UP QUERIES (CRITICAL)**: After successfully resolving an entity using search tools, you MUST incorporate the resolved entity ID/IDs into subsequent tool queries.
**Example Disambiguation Flow:**
- Query: "List work items assigned to Robert"
- Step 1: Call search_user_by_name("Robert") → Returns multiple users
- Step 2: Call ask_for_clarification with user options
- (System pauses for user input)
**EXECUTION STRATEGY:**
**Phase 1 - Initial Discovery (using the initial queries passed to you along with the selected tools):**
- **ENTITY DISAMBIGUATION FIRST**: If the query contains entity references, use entity search tools to resolve them before the main tools
- Use the EXACT query strings provided in the selected tools list for main tools **WITH ONE CRITICAL EXCEPTION**:
- **EXCEPTION: Entity ID Incorporation** - After successfully resolving entities (users, projects, cycles, modules, etc.) using search tools, you MUST replace the entity name/reference in the query with the resolved entity ID
- For example, if search_user_by_name returns user_id "abc-123", change "assigned to John" to "assigned to user with id: abc-123"
- This is NOT considered "modifying" the query - it's enriching it with resolved context
- Do not otherwise modify, rephrase, or optimize the provided tool queries
- Remember that these queries:
- are generated by me, not the user.
- They are designed to be as precise as possible.
- They are optimized for the specific tool's capabilities and limitations.
- Therefore, you should not modify them EXCEPT to incorporate resolved entity IDs.
- Don't pass text/string searches as part of structured_db_tool, especially if the text string is already assigned to a semantic search tool.
- Execute tools in logical order (disambiguation → discovery → enrichment)
- Empty results from any tool means you can skip to Phase 2.
**Phase 2 - Refinement (IF needed):**
- ONLY if Phase 1 results are insufficient to answer the user's question
- You may then modify queries to get additional details
- Clearly indicate this is a refinement attempt
**Example 1 - No Entity Resolution Needed:**
Given tools: [("vector_search_tool", "login issues"), ("structured_db_tool", "my tasks")]
Phase 1:
- Call vector_search_tool(query="login issues") # EXACT string
- Call structured_db_tool(query="my tasks") # EXACT string
**Example 2 - Entity Resolution Required:**
Given tools: [("structured_db_tool", "work items assigned to John")]
Phase 1:
- Call search_user_by_name("John") # Returns user_id: "abc-123-def"
- Call structured_db_tool(query="work items assigned to user with id: abc-123-def") # Incorporated resolved ID
Phase 2 (only if needed):
- If results insufficient, call tools with modified queries for more details
Execute Phase 1 first with exact queries (incorporating any resolved entity IDs). Only proceed to Phase 2 if necessary.
**Tool Usage Guidelines:**
**Query Formulation:**
- Use the exact queries as provided in the selected tools list, EXCEPT when incorporating resolved entity IDs (always replace entity names with resolved IDs).
- The structured_db_tool is a text2sql agent - it takes natural language input and converts it to SQL internally using Plane's database schema knowledge. Do NOT send SQL queries to this tool.
- Send natural language queries like "show me all high priority bugs assigned to John" or "count issues by state" to the structured_db_tool.
- **CRITICAL for structured_db_tool**: If you resolved entity IDs (from search_user_by_name, search_project_by_name, etc.), incorporate those IDs in the query. For example: "show me work items assigned to user with id: abc-123" instead of "assigned to John".
- Don't depend on the structured_db_tool to perform semantic search. Use the vector_search_tool for that.
- Only invoke vector_search_tool when the user specifies concrete content topics/keywords to match; do NOT trigger it for presentation-only instructions about transforming or displaying existing fields.
- If you collected issue_ids/page_ids from the semantic search, best to pass them in the IDs parameter of the structured_db_tool. Not in the query parameter.
- structured_db_tool: Send natural language queries (e.g., "show me all high priority bugs assigned to John"). It converts to SQL internally - do NOT send SQL.
- After resolving entity IDs: incorporate them in queries. Example: "work items assigned to user with id: abc-123" instead of "assigned to John".
- vector_search_tool: Use for semantic/content search on work-item titles and descriptions. Use structured_db_tool for filtering by metadata (assignees, states, dates).
- Pass semantic search results (issue_ids/page_ids) as IDs parameter to structured_db_tool, not in the query text.
**Note:** Internal tool parameters use 'issue_ids' (matching database schema), but always say "work-item" when communicating with users.
**Tool Selection Order:**
1. **Entity search first** if query mentions entity names (of any user, project, cycle, module, etc.)
2. **Choose primary retrieval tool:**
- Cycle metrics/progress → fetch_cycle_details (see cycle rules above)
- Semantic/content search → vector_search_tool or pages_search_tool
- Documentation questions → docs_search_tool
- Structured queries (assignments, states, counts) → structured_db_tool
3. **Enrichment** (if needed): Use additional tools to get more details
4. **Context building**: Use earlier tool outputs to inform later queries
5. **Skip dependent tools** if a tool returns no results
**CONTEXT-PROVIDED ENTITY IDS:**
When the user refers to entities using contextual references ("this project", "this cycle", "me", "my"), check if the relevant UUID is already available in the provided context:
- "this project" / "the project" → use project_id from context if available
- "this cycle" / "current cycle" → use cycle_id from context if available (or call search_current_cycle if not in context)
- "me" / "my" / "I" → use user_id from context
DO NOT call entity search tools (search_*_by_name) when the UUID is already available in context. Use the UUID directly in your tool calls.
Entity search tools should ONLY be called when:
- The user mentions a specific NAME (e.g., "project called Mobile", "user named David", "cycle named Sprint 5", "Workite type Bug")
- The required UUID is NOT available in the provided context
- You need to disambiguate between multiple entities with similar names
**ENTITY DISAMBIGUATION:**
When query mentions entity NAMES without IDs and the UUID is NOT in context:
1. Call appropriate search_*_by_name tool first (e.g., search_user_by_name, search_project_by_name)
2. If search returns multiple matches: call ask_for_clarification with the disambiguation options
3. If search returns no matches (e.g., "No user found", "Not found", "No project found", etc.):
- **YOU MUST call ask_for_clarification** - DO NOT just report the error to the user
- Use reason: "No [entity_type] found matching '[search_term]'"
- Use questions: ["Could you provide more details or check the spelling?", "Is there an alternative name or identifier?"]
- Use category_hints: [entity_type] (e.g., ["users"], ["projects"], ["workitems"])
- This is MANDATORY - never skip clarification when entity search fails
4. After resolving entities: incorporate resolved IDs in subsequent tool queries (e.g., "assigned to user with id: abc-123")
When IDs already provided: Skip entity search and use IDs directly.
Examples:
- Query "List work items assigned to Robert" + no user_id in context → search_user_by_name("Robert") → if multiple matches → ask_for_clarification
- Query "List my work items" + user_id in context → use user_id directly without calling entity search
- Query "is time tracking enabled in this project?" + project_id in context → use project_id directly without calling search_project_by_name
- Query "show me the Mobile project" + no project_id in context → search_project_by_name("Mobile") → use resolved project_id
**Default Incomplete Work Scope:**
- When the user's phrasing implies items that should be worked on now or are pending (e.g., "to pick up", "open", "pending", "to-do", "missed to add", "should be in this cycle", "not yet scheduled"), and the user does not explicitly ask for completed/closed items, include an uncompleted-state filter in the structured_db_tool query: state bucket/group IN (backlog, unstarted, started) and exclude completed/cancelled.
- Prefer filtering by the canonical state bucket/group if available; otherwise, use state names as a fallback in a case-insensitive manner.
- Combine this with any other constraints (e.g., membership, priority, assignees, dates) required by the user intent.
For queries about pending/open work (unless explicitly asking for completed items):
- Filter by state groups: backlog, unstarted, started (exclude completed/cancelled)
- Use state group field when available; otherwise use state names (case-insensitive)
- Combine with other constraints (assignees, priority, dates) as needed
**Priority Canonicalization:**
- Canonical priority values are: urgent, high, medium, low (case-insensitive). Do not invent other names like "highest".
- Map common synonyms to canonical values: "highest"/"critical"/"blocker"/"p0" → urgent; "very high"/"p1" → high; "normal"/"standard" → medium; "lowest" → low.
- When formulating structured_db_tool queries, use canonical values for priority filters and canonical urgent → high → medium → low order when sorting.
- For phrases like "highest priority" or "top priority" without an explicit priority name, do NOT equate to "urgent". Order by canonical priority and return the top item(s) without adding an equality filter unless the user explicitly says "urgent"/"high"/etc.
- Canonical values: urgent, high, medium, low (case-insensitive only)
- Map synonyms: "highest"/"critical"/"p0" → urgent, "very high"/"p1" → high, "normal" → medium, "lowest" → low
- Use canonical values in queries and for sorting
- For "top priority" / "highest priority": sort by priority and return top items (don't filter to urgent unless explicitly stated)
**WORKSPACE-LEVEL QUERIES (EFFICIENCY):**
|- When query is related to workitems and spans ALL projects without specifying a particular project or projects
| - DO NOT: call `list_member_projects` then query each project with `structured_db_tool` separately
| - INSTEAD: Use ONE `structured_db_tool` call WITHOUT `project_id` parameter
| - Query will automatically scope to workspace via `workspace_id` (already in context)
| - Example WRONG: `list_member_projects` + 7x `structured_db_tool` (each with different project_id)
| - Example RIGHT: Single `structured_db_tool`: "list all workitems prioritized as high"
|- For project-specific queries, include the specific `project_id`
**Project Features:**
- Cycles, modules, pages, workitem types, views, intake, and time tracking are project-level features that are enabled/disabled on a per-project basis.
- You can use:
- `projects_retrieve` tool to get details about whether the project features are enabled or disabled.
**Important:**
Only use the tools that were provided in the selected tools list, in your recommended execution order.
Only use the tools that were provided in the given tools list, in your recommended execution order.
Analyze the user's query and determine the most efficient sequential order that builds context progressively.
Always call tools in the logical order needed to answer the question completely.
When you have the complete answer, answer the user's question in a coherent and comprehensive manner.
Use the user's first name naturally in conversation when it feels appropriate.
Rules to follow while answering the user's question:
1. Ensure your answer directly addresses the user query.
3. **Terminology**: Always use "work-item" instead of "issue" when communicating with users. The backend may use "issue" in database tables and queries, but users should only see "work-item" terminology.
4. **Unique Keys**: Refer to work-item identifiers (like PAI-123, MOB-45) as "unique key" instead of "Issue ID" in user-facing responses.
5. Suppress the UUIDs (like User ID, Issue ID, Page ID, Project ID (The UUID of the project. Not the identifier), Workspace ID, etc) in your response. These are PII data. Never show them.
And don't mention the suppression in your response.
However, remember that when the user mentions 'issue ID' or 'issue identifier' he/she mean to refer to the issue identifier which is not UUIDs but the unique key like PAI-123, MOB-45.
For example, if the user asks "What is the issue ID of the issue 'Support for Custom Fields'?", and you are provided with the unique key 'PAI-123' in the responses by the tools, your response should be "The issue ID of the issue 'Support for Custom Fields' is PAI-123".
You can reveal the name of the user to him/her in your response, if requested.
In case of project, feel free to show the project identifier, but never show the UUID, even if the project_id (UUID) is mentioned by the user in the query. Only show the project identifier in your response.
6. If there is no data related to the user's query, just convey this. Don't hallucinate.
- IMPORTANT: Never mention internal details like "SQL", "query execution", or tool names when reporting no data.
- Phrase it in a user-friendly way tailored to the question, e.g., "No recent comments found on your work-items."
- Optionally, suggest one concise next step (e.g., broaden filters, time range) if helpful.
7. Never provide sensitive information: passwords, API keys, table names, schema details, or SQL queries.
8. If structured_db_tool has been deployed:
- List all items retrieved with relevant details (excluding UUIDs as per Rule 5)
- DO NOT re-filter results (already filtered by SQL query)
- Never reveal the SQL query to the end user
9. URL Embeddings: Create clickable links for entities with URLs in the Entity URLs section:
- Work-items: Entity Name [Unique Key](URL). Example: Support for Custom Fields [PAI-123](https://xyz.com/abc/123/)
- Other entities: Entity Name [view](URL). Example: Meeting Notes [view](https://xyz.com/abc/123/)
- In tables with separate columns: make the unique key/identifier column clickable [PAI-123](URL)
- Skip linking if entity not in Entity URLs section
10. When using tables to present data: Never include UUIDs (must be suppressed), and apply URL embedding rules from Rule 9.
""" # noqa: E501
# In pi/services/chat/prompts.py, add this constant:
RETRIEVAL_TOOL_DESCRIPTIONS = """
**Retrieval Tool Capabilities:**
1. **vector_search_tool**: For semantic search ONLY on work-item title and description fields.
- USE FOR: Finding work items by content, topics, keywords, concepts
- NOT FOR: Comments, updates, activity streams, state changes, metadata queries, presentation/output-formatting instructions
- RETURNS: Text results and a list of work-item IDs
2. **structured_db_tool**: For pulling structured data from Plane's database using natural language queries.
- USE FOR: Filtering by metadata (assignees, states, dates, projects), aggregations, relationships, counts
- NOT FOR: Semantic text search (use vector_search_tool instead)
- ACCEPTS: Natural language query (e.g., "show me all high priority bugs assigned to John") and optional issue_ids/page_ids from prior searches
- NOTE: This tool is a text2sql agent - it converts your natural language to SQL internally using Plane's database schema knowledge
3. **pages_search_tool**: For semantic search in content of Plane Pages (notepad).
- USE FOR: Finding pages by content, topics, concepts
- NOT FOR: Page metadata queries (who created, when created)
- RETURNS: Text results and a list of page IDs
4. **docs_search_tool**: For searching Plane's official documentation.
- USE FOR: Finding documentation by topics, features, how-to guides
- RETURNS: Formatted documentation search results
5. **generic_query_tool**: Fallback for non-Plane queries and pleasantries.
"""
apps/pi/pi/services/chat/templates.py
+39 -27
View File
@@ -50,8 +50,8 @@ def preset_question_flow(question: str) -> List[Dict[str, Any]]:
Returns a list of flow steps that define the complete processing pipeline.
Each step contains:
- step_type: The type of step (routing, agent_execution, etc.)
- agents: List of agents to execute with their queries and parameters
- step_type: The type of step (routing, tool_execution, etc.)
- agents: List of tools to execute with their queries and parameters (field name kept for backward compatibility)
- skip_routing: Whether to skip the normal routing process
- reasoning_messages: Messages to show during processing
"""
@@ -67,7 +67,7 @@ def preset_question_flow(question: str) -> List[Dict[str, Any]]:
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "Show me my urgent work items that are still pending",
"sql_query": """SELECT
issues.id::text AS issues_id,
@@ -99,7 +99,7 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "What work items are assigned to me and not yet completed, with start and end dates, priorities?",
"sql_query": """SELECT
issues.id::text AS issues_id,
@@ -135,7 +135,7 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "Show me my overdue work items",
"sql_query": """SELECT
issues.id::text AS issues_id,
@@ -148,7 +148,7 @@ LEFT JOIN states ON issues.state_id = states.id
WHERE
issue_assignees.assignee_id = $1
AND issues.target_date IS NOT NULL
AND issues.target_date < CURRENT_DATE
AND issues.target_date::date < CURRENT_DATE
AND (states.group != 'completed' AND states.group != 'cancelled')
AND issues.deleted_at IS NULL
AND issue_assignees.deleted_at IS NULL
@@ -170,7 +170,7 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "What did I complete this week?",
"sql_query": """SELECT
issues.id::text AS issues_id,
@@ -207,7 +207,7 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "Show me work items I'm working on currently",
"sql_query": """SELECT
issues.id::text AS issues_id,
@@ -242,7 +242,7 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "What work items are assigned to me that are blocked, and which ones are blocking them?",
"sql_query": """SELECT
i.id::text AS issues_id,
@@ -281,21 +281,33 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "Who commented on my work items recently?",
"sql_query": """SELECT
users.id::text AS users_id,
users.first_name,
users.last_name,
MAX(issue_comments.created_at) AS most_recent_comment_date
FROM issue_assignees
JOIN issues ON issue_assignees.issue_id = issues.id
JOIN issue_comments ON issue_comments.issue_id = issues.id
JOIN users ON users.id = issue_comments.created_by_id
WHERE issue_assignees.assignee_id = $1
AND issue_comments.deleted_at IS NULL
AND issue_assignees.deleted_at IS NULL
GROUP BY users.id, users.first_name, users.last_name
users_id,
first_name,
last_name,
most_recent_comment_date,
issues_id,
issue_name
FROM (
SELECT DISTINCT ON (users.id)
users.id::text AS users_id,
users.first_name,
users.last_name,
issue_comments.created_at AS most_recent_comment_date,
issues.id::text AS issues_id,
issues.name AS issue_name
FROM issue_assignees
JOIN issues ON issue_assignees.issue_id = issues.id
JOIN issue_comments ON issue_comments.issue_id = issues.id
JOIN users ON users.id = issue_comments.created_by_id
WHERE issue_assignees.assignee_id = $1
AND users.id <> $1
AND issue_comments.deleted_at IS NULL
AND issue_assignees.deleted_at IS NULL
ORDER BY users.id, issue_comments.created_at DESC
) t
ORDER BY most_recent_comment_date DESC
LIMIT 20;""",
"tables": ["issues", "users", "issue_assignees", "issue_comments"],
@@ -315,7 +327,7 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "What changed in my pending work items today?",
"sql_query": """SELECT
issues.id::text AS issues_id,
@@ -357,7 +369,7 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "Show me recent activity on my pending work items",
"sql_query": """SELECT
ia.id::text AS issue_activities_id,
@@ -406,7 +418,7 @@ LIMIT 20;""",
],
"agents": [
{
"agent": "plane_structured_database_agent",
"agent": "structured_db_tool",
"query": "What's the progress of current active cycles?",
"sql_query": """SELECT
cycles.id::text AS cycles_id,
@@ -463,7 +475,7 @@ ORDER BY cycles.start_date DESC NULLS LAST, cycles.name ASC;""",
# ]
# },
# {
# "step_type": "agent_execution",
# "step_type": "tool_execution",
# "reasoning_messages": [
# "🔧 Step 2: Retrieving detailed updates for the work item"
# ],
@@ -497,7 +509,7 @@ ORDER BY cycles.start_date DESC NULLS LAST, cycles.name ASC;""",
# ]
# },
# {
# "step_type": "agent_execution",
# "step_type": "tool_execution",
# "reasoning_messages": [
# "🔧 Step 2: Retrieving detailed updates for each work item"
# ],
apps/pi/pi/services/chat/utils.py
+106 -8
View File
@@ -1,4 +1,5 @@
import json
import random
import re
import uuid
from typing import Any
@@ -39,8 +40,6 @@ def mask_uuids_in_text(text: str) -> str:
Returns:
Cleaned text with UUIDs masked and context removed
"""
import re
uuid_pattern = r"\b[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\b"
def replace_uuid(match):
@@ -150,6 +149,22 @@ class StandardAgentResponse:
response_lower = response_text.lower()
return any(pattern.lower() in response_lower for pattern in error_patterns)
@staticmethod
def format_response_with_entity_urls(response: Union[str, Dict[str, Any]]) -> str:
"""Format a response into a string with entity URLs"""
if isinstance(response, dict):
entity_urls = response.get("entity_urls") or []
url_info = ""
if entity_urls:
url_info = "\n\n**Entity URLs Available:**\n"
for url in entity_urls:
url_info += f"- {url.get("type", "").title()}: {url.get("name")} - URL: {url.get("url")}\n"
if url.get("type") == "issue":
url_info += f"Issue Unique Key: {url.get("issue_identifier")}\n"
formatted_response = response.get("results", "") + url_info
return formatted_response
return str(response)
@staticmethod
def format_responses(responses, chat_id, query_flow_store):
formatted_responses = []
@@ -196,6 +211,89 @@ class StandardAgentResponse:
return {"formatted": formatted_responses_str, "has_errors": has_errors}
def reasoning_header_factory(stage: str, tool_name: str, tool_query: str) -> str:
"""Factory function to create a reasoning header for a given stage"""
stage_dict = {
"build_beginning": [
"🧠 Assembling the action sequence...\n\n",
"🤔 Planning next steps...\n\n",
"🧠 Brainstorming the best approach...\n\n",
"🛤️ Mapping out the action path...\n\n",
"🚧 Constructing your request workflow...\n\n",
"🔍 Evaluating available options...\n\n",
"⏳ Orchestrating the next steps...\n\n",
],
"selected_action_categories": [
"🧭 Identified planning focus\n\n",
"🔍 Evaluated relevant areas\n\n",
],
"actions_clarification_followup": [
"🧠 Continuing with your request...\n\n",
],
"planner_tool_selection_calling": [
"🧠 Zeroing in on the best tool...\n\n",
"🔍 Narrowing down the next moves...\n\n",
"🧭 Steps to take next...\n\n",
],
"planner_tool_selection": [
"💭 Contemplating...\n\n",
],
"planner_tool_selection_final": [
f"🤖 {tool_name} ({tool_query})\n\n" if tool_query else f"🤖 {tool_name}\n\n",
],
"retrieval_tool_execution": [
f"🔧 Executing: {tool_name}: {tool_query}\n\n",
f"Calling {tool_name} for {tool_query}\n\n",
],
"retrieval_tool_execution_message": [
f"✅Tool {tool_name} returned its output...\n\n",
f"I received {tool_name}'s result..\n\n",
],
"tool_complete": [
f"✅ {tool_name} execution completed\n\n",
f"{tool_name} is executed successfully\n\n",
],
"tool_error": [
f"⚠️ Error executing {tool_name}\n\n",
],
"tool_unavailable": [
"⚠️ A required capability was unavailable.\n\n",
],
"ask_mode_clarification_followup": [
"🤖 Processing your clarification...\n\n",
],
"ask_mode_beginning": [
"🤖 Retrieving information...\n\n",
"🤖 Processing your request...\n\n",
"🧠 Planning next steps...\n\n",
"🤔 Planning next steps...\n\n",
"🧠 Brainstorming the best approach...\n\n",
],
"ask_mode_analyzing_results": [
"🤖 Analyzing results...\n\n",
"🧠 Reviewing tool outputs...\n\n",
],
"ask_preset_reasoning": [
"🧭 Using optimized path...\n\n",
"💡 Applying optimized insights...\n\n",
"🧠 Following optimized guidance...\n\n",
],
"final_response": ["📝 Generating final response...\n\n"],
}
stage_list = stage_dict.get(stage, "🧠 {stage}\n\n")
return random.choice(stage_list)
def reasoning_dict_maker(stage: str, tool_name: str, tool_query: str, content: str) -> dict:
"""Factory function to create a reasoning dictionary for a given stage"""
return {
"chunk_type": "reasoning",
"header": mask_uuids_in_text(reasoning_header_factory(stage, tool_name, tool_query)),
"content": mask_uuids_in_text(content),
"stage": stage,
}
def format_conversation_history(conversation_history):
return "\n".join([f"{m.type}: {m.content}" for m in conversation_history])
@@ -257,12 +355,13 @@ async def process_conv_history(conv_history: list[dict[str, Any]], db: AsyncSess
# Extract and include action context if available
action_context = ""
execution_status = qa_pair.get("execution_status", {})
if execution_status.get("has_planned_actions") and execution_status.get("actions"):
# Fix: Remove has_planned_actions gate - just check if actions exist
if execution_status.get("actions"):
action_context += "\n**Action Context:**\n"
# Include planning context from actions
for action in execution_status.get("actions", []):
tool_name = action.get("tool")
action_type = action.get("action")
success = action.get("success")
if success and action.get("entity"):
@@ -271,11 +370,11 @@ async def process_conv_history(conv_history: list[dict[str, Any]], db: AsyncSess
entity_name = entity.get("entity_name", "")
entity_id = entity.get("entity_id", "")
action_context += f"- Action: {tool_name} → Successfully created/updated {entity_type} '{entity_name}' (ID: {entity_id})\n"
action_context += f"- Action: {action_type} → Successfully created/updated {entity_type} '{entity_name}' (ID: {entity_id})\n"
elif success:
action_context += f"- Action: {tool_name} → Executed successfully\n"
action_context += f"- Action: {action_type} → Executed successfully\n"
else:
action_context += f"- Action: {tool_name} → Failed to execute\n"
action_context += f"- Action: {action_type} → Failed to execute\n"
# Check if we can extract planning context from the qa_pair structure
# This would be available when the enhanced execution_data includes planning_context
@@ -441,7 +540,6 @@ async def initialize_new_chat(
chat_exists = await check_if_chat_exists(final_chat_id, db)
if chat_exists:
return {"success": False, "message": "Chat ID already exists", "error_code": "CHAT_EXISTS"}
log.info(f"ChatID: {final_chat_id} - Initializing chat context to db: is_project_chat={is_project_chat}")
# Create the chat record without workspace details (will be backfilled later)
chat_result = await upsert_chat(
chat_id=final_chat_id,
apps/pi/pi/services/dupes/dupes.py
+3 -3
View File
@@ -1,7 +1,7 @@
import time
from langchain.schema import HumanMessage
from langchain.schema import SystemMessage
from langchain_core.messages import HumanMessage
from langchain_core.messages import SystemMessage
from pi import logger
from pi import settings
@@ -83,7 +83,7 @@ async def identify_duplicates_with_llm(query_title: str, query_description: str,
# Create the dupes LLM with structured output (include_raw=True to preserve token usage)
dupes_llm = get_dupes_llm()
dupes_structured_llm = dupes_llm.with_structured_output(DuplicateIdentificationResponse, include_raw=True) # type: ignore[arg-type]
dupes_structured_llm = dupes_llm.with_structured_output(DuplicateIdentificationResponse, include_raw=True, method="json_schema") # type: ignore[arg-type]
# Format candidates for LLM input as a clean numbered list
candidates_text = ""
apps/pi/pi/services/llm/llms.py
+51 -7
View File
@@ -11,7 +11,7 @@ from typing import Iterator
from typing import Optional
from uuid import UUID
from langchain.schema.language_model import BaseLanguageModel
from langchain_core.language_models.base import BaseLanguageModel
from langchain_core.runnables import Runnable
from langchain_core.runnables import RunnableConfig
from langchain_openai import ChatOpenAI
@@ -592,9 +592,53 @@ def get_dupes_llm() -> BaseLanguageModel:
return create_openai_llm(fallback_config)
# Global instances for backward compatibility
llm = LLMFactory.get_default_llm()
stream_llm = LLMFactory.get_stream_llm()
decomposer_llm = LLMFactory.get_decomposer_llm()
fast_llm = LLMFactory.get_fast_llm(streaming=False)
fast_llm_stream = LLMFactory.get_fast_llm(streaming=True)
class LazyLLM(Runnable):
"""Lazy initialization proxy for LLMs to avoid shared SSL contexts in forked processes."""
def __init__(self, factory_func: Any):
self._factory_func = factory_func
self._proxy_target: Any = None
@property
def _target(self) -> Any:
if self._proxy_target is None:
self._proxy_target = self._factory_func()
return self._proxy_target
def invoke(self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any) -> Any: # noqa: A002
return self._target.invoke(input, config, **kwargs)
async def ainvoke(self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any) -> Any: # noqa: A002
return await self._target.ainvoke(input, config, **kwargs)
def stream(self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any) -> Iterator[Any]: # noqa: A002
return self._target.stream(input, config, **kwargs)
async def astream(self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any) -> AsyncIterator[Any]: # noqa: A002
async for chunk in self._target.astream(input, config, **kwargs):
yield chunk
def bind_tools(self, *args: Any, **kwargs: Any) -> Any:
return self._target.bind_tools(*args, **kwargs)
def with_structured_output(self, *args: Any, **kwargs: Any) -> Any:
return self._target.with_structured_output(*args, **kwargs)
def __getattr__(self, name: str) -> Any:
return getattr(self._target, name)
@property
def InputType(self) -> Any:
return self._target.InputType
@property
def OutputType(self) -> Any:
return self._target.OutputType
# Global instances for backward compatibility - using LazyLLM to prevent fork issues
llm = LazyLLM(LLMFactory.get_default_llm)
stream_llm = LazyLLM(LLMFactory.get_stream_llm)
decomposer_llm = LazyLLM(LLMFactory.get_decomposer_llm)
fast_llm = LazyLLM(lambda: LLMFactory.get_fast_llm(streaming=False))
fast_llm_stream = LazyLLM(lambda: LLMFactory.get_fast_llm(streaming=True))
apps/pi/pi/services/retrievers/docs_search.py
+2 -2
View File
@@ -1,5 +1,5 @@
from langchain.callbacks.manager import AsyncCallbackManagerForRetrieverRun
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain_core.callbacks.manager import AsyncCallbackManagerForRetrieverRun
from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
apps/pi/pi/services/retrievers/issue_search.py
+4 -2
View File
@@ -1,7 +1,7 @@
import asyncio
from langchain.callbacks.manager import AsyncCallbackManagerForRetrieverRun
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain_core.callbacks.manager import AsyncCallbackManagerForRetrieverRun
from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
@@ -96,6 +96,8 @@ class IssueRetriever(BaseRetriever):
Document(
page_content=description,
metadata={
# Provide both for compatibility with downstream formatters
"title": title,
"name": title,
"relevance": round(hit.get("Score", 0), 3),
"issue_id": issue_id,
apps/pi/pi/services/retrievers/pages_search.py
+4 -3
View File
@@ -1,7 +1,7 @@
from typing import List
from langchain.callbacks.manager import AsyncCallbackManagerForRetrieverRun
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain_core.callbacks.manager import AsyncCallbackManagerForRetrieverRun
from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
@@ -76,7 +76,8 @@ class PageChunkRetriever(BaseRetriever):
Document(
page_content=description,
metadata={
"page_title": title,
# Use 'name' to match downstream formatter
"name": title,
"relevance": round(hit.get("Score", 0), 3),
"page_id": issue_id,
},
apps/pi/pi/services/retrievers/pg_store/action_artifact.py
+68 -5
View File
@@ -1,3 +1,4 @@
from datetime import datetime
from typing import Any
from typing import Dict
from typing import List
@@ -12,6 +13,7 @@ from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.app.models import ActionArtifact
from pi.app.models.action_artifact import ActionArtifactVersion
from pi.app.models.enums import ExecutionStatus
from pi.app.models.enums import FlowStepType
from pi.app.models.message import MessageFlowStep
@@ -35,7 +37,6 @@ async def get_action_artifacts_by_ids(
.where(ActionArtifact.deleted_at.is_(None)) # type: ignore[union-attr,arg-type]
.order_by(desc(ActionArtifact.created_at)) # type: ignore[union-attr,arg-type]
)
result = await db.execute(stmt)
artifacts = list(result.scalars().all())
@@ -148,19 +149,27 @@ async def create_action_artifact(
async def update_action_artifact_execution_status(
db: AsyncSession,
artifact_id: UUID4,
message_id: UUID4,
chat_id: UUID4,
is_executed: bool,
success: bool,
entity_id: Optional[UUID4] = None,
entity_info: Optional[Dict[str, Any]] = None,
execution_result: Optional[str] = None,
executed_at: Optional[str] = None,
) -> Dict[str, Any]:
"""
Updates the execution status of an action artifact.
Updates the execution status of an action artifact and corresponding MessageFlowStep.
Args:
db: Database session
artifact_id: ID of the artifact to update
is_executed: New execution status
success: Whether the execution succeeded
entity_id: Optional entity ID to set (if entity was created)
entity_info: Optional entity info dict (entity_url, entity_name, etc.)
execution_result: Optional execution result message
executed_at: Optional execution timestamp (ISO format string)
Returns:
A dictionary with operation status and the artifact object or error details.
@@ -194,8 +203,62 @@ async def update_action_artifact_execution_status(
# Don't fail if data shaping has issues; persistence of status is primary
pass
# Add and commit
# Add artifact to session
db.add(artifact)
# Also update MessageFlowStep if artifact has a message_id
if message_id:
try:
# Find the MessageFlowStep for this artifact
flow_step_stmt = (
select(MessageFlowStep)
.where(MessageFlowStep.chat_id == chat_id) # type: ignore[arg-type]
.where(MessageFlowStep.message_id == message_id) # type: ignore[arg-type]
.where(MessageFlowStep.execution_data.op("->>")("artifact_id") == str(artifact_id)) # type: ignore[union-attr,arg-type]
.where(MessageFlowStep.is_planned == True) # type: ignore[arg-type] # noqa: E712
)
flow_step_result = await db.execute(flow_step_stmt)
flow_step = flow_step_result.scalar_one_or_none()
if flow_step:
# Mark as executed
flow_step.is_executed = is_executed
# Set execution success status
if success:
flow_step.execution_success = ExecutionStatus.SUCCESS
flow_step.execution_error = None
else:
flow_step.execution_success = ExecutionStatus.FAILED
error_msg = (
execution_result
if execution_result and "error" in execution_result.lower()
else f"Execution failed: {execution_result or "Unknown error"}"
)
flow_step.execution_error = error_msg
# Update execution_data
current_execution_data = flow_step.execution_data or {}
updated_execution_data = {
**current_execution_data,
"executed_at": executed_at or datetime.utcnow().isoformat(),
"execution_result": execution_result or "",
"business_success": success,
}
if entity_info:
updated_execution_data["entity_info"] = entity_info
flow_step.execution_data = updated_execution_data
attributes.flag_modified(flow_step, "execution_data")
db.add(flow_step)
else:
log.warning(f"MessageFlowStep not found for artifact {artifact_id} and message {message_id}")
except Exception as flow_step_error:
log.error(f"Error updating MessageFlowStep for artifact {artifact_id}: {flow_step_error}")
# Don't fail the main execution if flow step update fails
# Commit all changes
await db.commit()
await db.refresh(artifact)
@@ -528,7 +591,7 @@ async def add_query_to_artifact(db: AsyncSession, artifact_id: UUID4, message_id
select(MessageFlowStep)
.where(
MessageFlowStep.step_type == "ARTIFACT_CHAT", # type: ignore[union-attr,arg-type]
MessageFlowStep.execution_data.op("->>")('"artifact_id"') == str(artifact_id), # type: ignore[union-attr,arg-type]
MessageFlowStep.execution_data.op("->>")("artifact_id") == str(artifact_id), # type: ignore[union-attr,arg-type]
MessageFlowStep.message_id == message_id, # type: ignore[union-attr,arg-type]
)
.order_by(desc(MessageFlowStep.created_at)) # type: ignore[union-attr,arg-type]
@@ -576,7 +639,7 @@ async def get_artifact_prompt_history_from_flow_steps(db: AsyncSession, artifact
try:
conditions = [
MessageFlowStep.step_type == FlowStepType.ARTIFACT_CHAT, # type: ignore[union-attr,arg-type]
MessageFlowStep.execution_data.op("->>")('"artifact_id"') == str(artifact_id), # type: ignore[union-attr,arg-type]
MessageFlowStep.execution_data.op("->>")("artifact_id") == str(artifact_id), # type: ignore[union-attr,arg-type]
MessageFlowStep.message_id != message_id, # type: ignore[union-attr,arg-type]
]
apps/pi/pi/services/retrievers/pg_store/chat.py
+386 -104
View File
@@ -3,6 +3,7 @@ import uuid
from typing import Any
from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from typing import Tuple
from typing import Union
@@ -14,11 +15,13 @@ from sqlalchemy import select
from sqlmodel.ext.asyncio.session import AsyncSession
from pi import logger
from pi.agents.sql_agent.tools import format_as_bullet_points
from pi.app.models import Chat
from pi.app.models import Message
from pi.app.models import MessageFlowStep
from pi.app.models import UserChatPreference
from pi.app.models.enums import ExecutionStatus
from pi.app.models.enums import FocusEntityType
from pi.app.models.enums import UserTypeChoices
from pi.app.models.message_attachment import MessageAttachment
from pi.app.schemas.chat import PaginationResponse
@@ -27,6 +30,8 @@ from pi.app.utils.attachments import get_presigned_url_preview
from pi.app.utils.pagination import apply_cursor_pagination
from pi.app.utils.pagination import check_pagination_bounds
from pi.app.utils.pagination import create_pagination_response
from pi.services.chat.helpers.tool_utils import format_tool_message_for_display
from pi.services.chat.utils import mask_uuids_in_text
from pi.services.query_utils import parse_query
log = logger.getChild(__name__)
@@ -35,10 +40,10 @@ log = logger.getChild(__name__)
internal_reasoning_format_dict = {
"rewrite": "Rewritten Query",
"routing": "Routing",
"tool": "Selected Agent",
"PLANE_STRUCTURED_DATABASE_AGENT": "Plane Structured Database Agent",
"PLANE_VECTOR_SEARCH_AGENT": "Plane Vector Search Agent",
"PLANE_PAGES_AGENT": "Plane Pages Search Agent",
"tool": "Selected Tool",
"structured_db_tool": "Structured DB Tool",
"vector_search_tool": "Vector Search Tool",
"pages_search_tool": "Pages Search Tool",
}
@@ -502,6 +507,171 @@ async def retrieve_chat_history(
"execution_status": execution_status_info,
}
# For external responses, augment reasoning with cleaned tool results from flow steps
if not pi_internal:
try:
# Collect relevant flow steps for this message, in order
message_specific_steps = [step for step in message_flow_steps if step.message_id == user_message.id]
message_specific_steps.sort(key=lambda x: x.step_order or 0)
cleaned_blocks: List[str] = []
for step in message_specific_steps:
# Only include tool steps; skip metadata/internal-only tools
if getattr(step, "step_type", None) != "tool":
continue
tool_name = getattr(step, "tool_name", None)
if tool_name in (
"tool_orchestration_context",
"llm_reasoning",
"get_available_plane_actions",
"planner_summary",
):
continue
# Skip planned actions in external reasoning (only show retrieval outputs)
if getattr(step, "is_planned", False):
continue
# Include retrieval steps even if execution_success wasn't explicitly set,
# as ask-mode persistence used to leave it as PENDING. Prefer executed or explicit success.
if (
not getattr(step, "is_executed", False)
and getattr(step, "execution_success", None) != ExecutionStatus.SUCCESS
):
continue
# SPECIAL FILTERS:
# For structured_db_tool, hide intermediate steps (table selection, SQL generation, final_query planning)
if isinstance(tool_name, str) and tool_name.startswith("structured_db_tool_"):
# Only the core 'structured_db_tool' should be shown in external reasoning
continue
# Extract raw content and format for display
exec_data = getattr(step, "execution_data", {}) or {}
raw_str: str = ""
if isinstance(exec_data, dict):
if exec_data.get("execution_result"):
raw_str = str(exec_data.get("execution_result"))
elif exec_data.get("retrieval_result"):
raw_str = str(exec_data.get("retrieval_result"))
elif exec_data.get("structured_result") is not None:
try:
raw_str = json.dumps(exec_data.get("structured_result"), default=str)
except Exception:
raw_str = str(exec_data.get("structured_result"))
if not raw_str:
raw_content = parse_flow_step_content(step.content or "")
try:
if isinstance(raw_content, (dict, list)):
raw_str = json.dumps(raw_content, default=str)
else:
raw_str = str(raw_content)
except Exception:
raw_str = str(raw_content)
# Tool-specific external rendering
tname = (tool_name or "").strip().lower()
if tname == "structured_db_tool":
try:
# Prefer structured_result.results if available
results_obj = None
sql_query_str = None
if isinstance(exec_data, dict):
sr = exec_data.get("structured_result")
if isinstance(sr, dict):
results_obj = sr.get("results")
sql_query_str = sr.get("sql_query")
if results_obj is None and raw_str:
# Try to parse raw_str JSON and extract 'results'
try:
parsed = json.loads(raw_str)
if isinstance(parsed, dict):
results_obj = parsed.get("results")
if not sql_query_str:
sql_query_str = parsed.get("sql_query")
except Exception:
results_obj = None
# Fallback: if still missing, try step.content JSON
if results_obj is None:
content_parsed = parse_flow_step_content(step.content or "")
if isinstance(content_parsed, dict):
results_obj = content_parsed.get("results")
if not sql_query_str:
sql_query_str = content_parsed.get("sql_query")
# Build display exactly as live stream intent: success prelude + rows
prelude = "✅ Database querying execution completed"
rows_text = ""
if results_obj is not None:
# Use bullet points formatter (async)
rows_text = await format_as_bullet_points(results_obj, sql_query=sql_query_str)
else:
# Last resort: fall back to cleaned message formatter
rows_text = format_tool_message_for_display(f"{prelude}\n\nResult: {raw_str}")
# Avoid duplicating prelude inside rows_text
if rows_text.startswith(prelude):
rows_text = rows_text[len(prelude) :].lstrip()
block = f"{prelude}\n\n{rows_text}".strip()
if block:
cleaned_blocks.append(mask_uuids_in_text(block))
except Exception:
# Fail-safe: fall back to generic formatting
prelude = f"✅ {(tool_name or "Tool")} execution completed"
formatted = format_tool_message_for_display(f"{prelude}\n\nResult: {raw_str}")
if formatted and formatted.strip():
cleaned_blocks.append(mask_uuids_in_text(formatted))
elif tname == "vector_search_tool":
# Truncate to first 50 words for external display
def _truncate_words(text: str, max_words: int = 50) -> str:
words = text.split()
if len(words) <= max_words:
return text
return " ".join(words[:max_words]) + "..."
truncated = ""
# Try to extract 'results' from structured_result if present
if isinstance(exec_data, dict) and exec_data.get("structured_result"):
sr = exec_data.get("structured_result")
if isinstance(sr, dict):
res = sr.get("results")
if res is not None:
truncated = _truncate_words(str(res))
if not truncated:
truncated = _truncate_words(str(raw_str))
from pi.services.chat.helpers.tool_utils import tool_name_shown_to_user
prelude = f"✅ {tool_name_shown_to_user(tool_name or "Tool")} execution completed"
formatted = f"{prelude}\n\nResult: {truncated}"
cleaned_blocks.append(mask_uuids_in_text(formatted.strip()))
else:
from pi.services.chat.helpers.tool_utils import tool_name_shown_to_user
prelude = f"✅ {tool_name_shown_to_user(tool_name or "Tool")} execution completed"
# Include query context when available (helps external readers)
if isinstance(exec_data, dict) and exec_data.get("tool_query"):
try:
tq = str(exec_data.get("tool_query")).strip()
if tq:
prelude = f"{prelude} ({tq})"
except Exception:
pass
formatted = format_tool_message_for_display(f"{prelude}\n\nResult: {raw_str}")
if formatted and formatted.strip():
cleaned_blocks.append(mask_uuids_in_text(formatted))
if cleaned_blocks:
base_reasoning = (qa_pair.get("reasoning") or "").strip()
if base_reasoning:
base_reasoning = mask_uuids_in_text(base_reasoning)
appended = "\n\n".join(cleaned_blocks)
qa_pair["reasoning"] = f"{base_reasoning}\n\n{appended}".strip() if base_reasoning else appended
except Exception:
# Fail-safe: do not block retrieval on formatting errors
pass
# Add execution information at dialogue level (only if there are actions)
if execution_status_info.get("actions"):
qa_pair.update({
@@ -523,57 +693,57 @@ async def retrieve_chat_history(
execution_summaries: List[str] = []
# Extract information from flow steps for this message
selected_agents_results: List[str] = []
agent_results: Dict[str, List[Any]] = {}
selected_tools_results: List[str] = []
tool_results: Dict[str, List[Any]] = {}
# Process flow steps to extract information
for message_flow_step in message_specific_steps:
content = parse_flow_step_content(message_flow_step.content or "")
step_type = message_flow_step.step_type
tool_name = message_flow_step.tool_name
if step_type == "routing":
# Handle routing step - content could be a dict with a list or a string
routing_data = None
if isinstance(content, dict):
# Check if content contains a list of routing results
routing_data = content.get("results")
elif isinstance(content, list):
routing_data = content
elif isinstance(content, str):
routing_data = content
if routing_data is not None:
if isinstance(routing_data, list):
# Content contains the list of selected agents
routing_results = []
for d in routing_data:
if not isinstance(d, dict):
continue
tool_val = str(d.get("tool", "")).strip()
query_val = str(d.get("query", "")).strip()
if not tool_val or not query_val:
# Filter empty agent entries safely
continue
tool_disp = internal_reasoning_format_dict.get(tool_val, tool_val)
routing_results.append(
f"Selected Agent: {tool_disp} and the corresponding sub query: {query_val}"
)
if routing_results:
selected_agents_results.extend(routing_results)
elif isinstance(routing_data, str):
selected_agents_results.append(routing_data)
else:
log.warning(f"Unexpected routing data type: {type(routing_data)}, routing_data: {routing_data}")
elif step_type == "tool" and tool_name:
# Skip context-only metadata steps (don't clutter history)
if tool_name == "tool_orchestration_context":
if step_type == "tool" and tool_name:
# Handle tool selection specially - extract selected tools for display
if tool_name == "tool_selection":
content_data = parse_flow_step_content(message_flow_step.content or "")
if isinstance(content_data, dict):
selected_tools = content_data.get("selected_tools", [])
for tool_info in selected_tools:
if isinstance(tool_info, dict):
tool_display_name = tool_info.get("name", "")
# Format tool selection for enhanced history
selected_tools_results.append(f"Selected Tool: {tool_display_name}")
continue
# Collect agent results
if tool_name not in agent_results:
agent_results[tool_name] = []
agent_results[tool_name].append(content)
# Skip context-only metadata steps, verbose advisory content, and reasoning (don't clutter history)
if tool_name in (
"tool_orchestration_context",
"llm_reasoning",
"get_available_plane_actions",
"planner_summary",
):
continue
# Skip intermediate retrieval tool steps (same as external reasoning)
# These are internal implementation details that shouldn't clutter the LLM context
if tool_name in (
"structured_db_tool_table_selection",
"structured_db_tool_sql_generation",
"structured_db_tool_final_query",
):
continue
# For retrieval tools (non-planned), skip PENDING steps - they shouldn't be in history
if hasattr(message_flow_step, "execution_success") and hasattr(message_flow_step, "is_planned"):
if not message_flow_step.is_planned:
# This is a retrieval tool - only include if successfully executed
if message_flow_step.execution_success != ExecutionStatus.SUCCESS:
continue
# Collect tool results
if tool_name not in tool_results:
tool_results[tool_name] = []
tool_results[tool_name].append(content)
# Add execution status information using new fields
if hasattr(message_flow_step, "execution_success") and hasattr(message_flow_step, "is_planned"):
@@ -735,11 +905,11 @@ async def retrieve_chat_history(
# Prepare separate sections so downstream can reorder (selected vs executed)
selected_section_lines: List[str] = []
if selected_agents_results:
if selected_tools_results:
selected_section_lines.append(
"The below agents were selected to retrieve relevant information to address the query:"
"The below tools were selected to retrieve relevant information to address the query:"
)
selected_section_lines.extend(selected_agents_results)
selected_section_lines.extend(selected_tools_results)
selected_section_lines.append("")
# Also add into internal_reasoning_parts for backwards-compatible single string
internal_reasoning_parts.extend(selected_section_lines)
@@ -752,12 +922,12 @@ async def retrieve_chat_history(
# Also add into internal_reasoning_parts for backwards-compatible single string
internal_reasoning_parts.extend(executed_section_lines)
# Process each agent's results - add only non-redundant, high-signal details
# Process each tool's results - add only non-redundant, high-signal details
# Avoid duplicating raw results here since execution summaries above already include a results preview
for agent_name, agent_content_list in agent_results.items():
for content in agent_content_list:
for tool_name_key, tool_content_list in tool_results.items():
for content in tool_content_list:
if isinstance(content, dict):
if agent_name == "PLANE_STRUCTURED_DATABASE_AGENT":
if tool_name_key == "structured_db_tool":
# Include generated SQL query (useful context), but skip raw result dumps
intermediate_results = content.get("intermediate_results", {})
if intermediate_results:
@@ -778,20 +948,29 @@ async def retrieve_chat_history(
f" issue unique key: {url_dict.get("issue_identifier")}"
)
else:
# Skip adding additional raw results for other agents to avoid duplication
# Skip adding additional raw results for other tools to avoid duplication
continue
internal_reasoning_parts.append("") # Add empty line between agents
internal_reasoning_parts.append("") # Add empty line between tools
# Join all internal reasoning parts into a single formatted string
message_internal_reasoning = "\n".join(internal_reasoning_parts)
# Store both combined and split sections for downstream formatting control
qa_pair["internal_reasoning"] = message_internal_reasoning
if selected_section_lines:
qa_pair["internal_selected"] = "\n".join(selected_section_lines).strip()
if executed_section_lines:
qa_pair["internal_executed"] = "\n".join(executed_section_lines).strip()
try:
if selected_tools_results:
qa_pair["internal_selected"] = "\n".join(
["The below tools were selected to retrieve relevant information to address the query:"]
+ selected_tools_results
).strip()
except Exception:
pass
try:
if executed_section_lines:
qa_pair["internal_executed"] = "\n".join(executed_section_lines).strip()
except Exception:
pass
dialogue_list.append(qa_pair)
i += 2
@@ -821,10 +1000,18 @@ async def retrieve_chat_history(
"query_id": str(user_message.id),
"answer_id": "", # No assistant message yet
"attachment_ids": attachments_dict.get(user_message.id, []),
"attachments": attachments_object_dict.get(user_message.id, []), # Add this line for consistency
# Add execution status information for frontend
"execution_status": execution_status_info,
}
# Flatten actions and summary at the top-level for frontend compatibility (build mode)
if execution_status_info.get("actions"):
standalone_qa_pair.update({
"action_summary": execution_status_info.get("action_summary", {}),
"actions": execution_status_info.get("actions", []),
})
if pi_internal:
# Generate internal reasoning for this standalone user message
standalone_message_internal_reasoning = ""
@@ -840,8 +1027,8 @@ async def retrieve_chat_history(
# Extract information from flow steps for this message
standalone_original_query = user_message.content or ""
standalone_rewritten_query = ""
standalone_selected_agents_results: List[str] = []
standalone_agent_results: Dict[str, List[Any]] = {}
standalone_selected_tools_results: List[str] = []
standalone_tool_results: Dict[str, List[Any]] = {}
# Process flow steps to extract information
for message_flow_step in standalone_message_specific_steps:
@@ -851,36 +1038,48 @@ async def retrieve_chat_history(
if step_type == "rewrite" and isinstance(content, dict):
standalone_rewritten_query = content.get("results", "") or content.get("rewritten_query", "")
elif step_type == "routing":
# Handle routing step - content could be a dict with a list or a string
routing_data = None
if isinstance(content, dict):
# Check if content contains a list of routing results
routing_data = content.get("results")
elif isinstance(content, list):
routing_data = content
elif isinstance(content, str):
routing_data = content
if routing_data is not None:
if isinstance(routing_data, list):
# Content contains the list of selected agents
routing_results = [
f"Selected Agent: {internal_reasoning_format_dict.get(d.get("tool", ""), d.get("tool", ""))} and the corresponding sub query: {d.get("query", "")}" # noqa: E501
for d in routing_data
if isinstance(d, dict)
]
standalone_selected_agents_results.extend(routing_results)
elif isinstance(routing_data, str):
standalone_selected_agents_results.append(routing_data)
else:
log.warning(f"Unexpected routing data type: {type(routing_data)}, routing_data: {routing_data}")
elif step_type == "tool" and tool_name:
# Collect agent results
if tool_name not in standalone_agent_results:
standalone_agent_results[tool_name] = []
standalone_agent_results[tool_name].append(content)
# Handle tool selection specially - extract selected tools for display
if tool_name == "tool_selection":
content_data = parse_flow_step_content(message_flow_step.content or "")
if isinstance(content_data, dict):
selected_tools = content_data.get("selected_tools", [])
for tool_info in selected_tools:
if isinstance(tool_info, dict):
tool_display_name = tool_info.get("name", "")
standalone_selected_tools_results.append(f"Selected Tool: {tool_display_name}")
continue
# Skip context-only metadata steps, verbose advisory content, and reasoning
if tool_name in (
"tool_orchestration_context",
"llm_reasoning",
"get_available_plane_actions",
"planner_summary",
):
continue
# Skip intermediate retrieval tool steps (same as external reasoning)
# These are internal implementation details that shouldn't clutter the LLM context
if tool_name in (
"structured_db_tool_table_selection",
"structured_db_tool_sql_generation",
"structured_db_tool_final_query",
):
continue
# For retrieval tools (non-planned), skip PENDING steps - they shouldn't be in history
if hasattr(message_flow_step, "execution_success") and hasattr(message_flow_step, "is_planned"):
if not message_flow_step.is_planned:
# This is a retrieval tool - only include if successfully executed
if message_flow_step.execution_success != ExecutionStatus.SUCCESS:
continue
# Collect tool results
if tool_name not in standalone_tool_results:
standalone_tool_results[tool_name] = []
standalone_tool_results[tool_name].append(content)
# Add execution status information using new fields
if hasattr(message_flow_step, "execution_success") and hasattr(message_flow_step, "is_planned"):
@@ -939,19 +1138,19 @@ async def retrieve_chat_history(
standalone_internal_reasoning_parts.append(f"User question: '{cleaned_query}'")
standalone_internal_reasoning_parts.append("")
if standalone_selected_agents_results:
if standalone_selected_tools_results:
standalone_internal_reasoning_parts.append(
"The below agents were selected to retrieve relevant information to address the query:"
"The below tools were selected to retrieve relevant information to address the query:"
)
standalone_internal_reasoning_parts.extend(standalone_selected_agents_results)
standalone_internal_reasoning_parts.extend(standalone_selected_tools_results)
standalone_internal_reasoning_parts.append("")
# Process each agent's results
for agent_name, agent_content_list in standalone_agent_results.items():
for content in agent_content_list:
# Process each tool's results
for tool_name_key, tool_content_list in standalone_tool_results.items():
for content in tool_content_list:
if isinstance(content, dict):
# Handle structured database agent specifically
if agent_name == "PLANE_STRUCTURED_DATABASE_AGENT":
# Handle structured database tool specifically
if tool_name_key == "structured_db_tool":
# standalone_internal_reasoning_parts.append("Intermediate results")
# SQL query and results
@@ -978,8 +1177,8 @@ async def retrieve_chat_history(
f" issue unique key: {url_dict.get("issue_identifier")}"
)
# Handle vector search agent
elif agent_name == "PLANE_VECTOR_SEARCH_AGENT":
# Handle vector search tool
elif tool_name_key == "vector_search_tool":
results = content.get("results", "")
if results:
standalone_internal_reasoning_parts.append("Semantic search results:")
@@ -999,8 +1198,8 @@ async def retrieve_chat_history(
standalone_internal_reasoning_parts.append(
f" issue unique key: {url_dict.get("issue_identifier")}"
)
# Handle pages search agent
elif agent_name == "PLANE_PAGES_AGENT":
# Handle pages search tool
elif tool_name_key == "pages_search_tool":
results = content.get("results", "")
if results:
standalone_internal_reasoning_parts.append("Pages search results:")
@@ -1008,7 +1207,7 @@ async def retrieve_chat_history(
else:
standalone_internal_reasoning_parts.append("Pages search results: No results found")
# Generic handling for other agents
# Generic handling for other tools
else:
results = content.get("results", "")
if results:
@@ -1019,7 +1218,7 @@ async def retrieve_chat_history(
standalone_internal_reasoning_parts.append("Results:")
standalone_internal_reasoning_parts.append(content)
standalone_internal_reasoning_parts.append("") # Add empty line between agents
standalone_internal_reasoning_parts.append("") # Add empty line between tools
# Join all internal reasoning parts into a single formatted string
standalone_message_internal_reasoning = "\n".join(standalone_internal_reasoning_parts)
@@ -1031,7 +1230,8 @@ async def retrieve_chat_history(
else:
i += 1
return {
# Return both new polymorphic structure and legacy fields for backward compatibility
response = {
"chat_id": str(chat_id),
"title": chat.title or "",
"dialogue": dialogue_list,
@@ -1039,13 +1239,19 @@ async def retrieve_chat_history(
"reasoning": "",
"llm": chat_llm,
"is_focus_enabled": user_chat_preference.is_focus_enabled if user_chat_preference else False,
# New polymorphic structure
"focus_entity_type": user_chat_preference.focus_entity_type if user_chat_preference and user_chat_preference.focus_entity_type else None,
"focus_entity_id": str(user_chat_preference.focus_entity_id) if user_chat_preference and user_chat_preference.focus_entity_id else None,
# Legacy fields (for backward compatibility)
"focus_project_id": str(user_chat_preference.focus_project_id)
if user_chat_preference and user_chat_preference.focus_project_id
else None,
"focus_workspace_id": str(user_chat_preference.focus_workspace_id)
if user_chat_preference and user_chat_preference.focus_workspace_id
else None,
"mode": user_chat_preference.mode if user_chat_preference and user_chat_preference.mode else "ask",
}
return response
except Exception as e:
log.error(f"Error retrieving chat history: {e}")
@@ -1057,8 +1263,11 @@ async def retrieve_chat_history(
"reasoning": "",
"llm": "",
"is_focus_enabled": False,
"focus_entity_type": None,
"focus_entity_id": None,
"focus_project_id": None,
"focus_workspace_id": None,
"mode": "ask",
}
@@ -1495,25 +1704,54 @@ async def upsert_user_chat_preference(
chat_id: UUID4,
db: AsyncSession,
is_focus_enabled: Optional[bool] = None,
# New polymorphic parameters
focus_entity_type: Optional[str] = None,
focus_entity_id: Optional[UUID4] = None,
# Legacy parameters (for backward compatibility)
focus_project_id: Optional[UUID4] = None,
focus_workspace_id: Optional[UUID4] = None,
mode: Optional[Literal["ask", "build"]] = "ask",
) -> Dict[str, Any]:
"""
Upserts a user chat preference.
Returns a dictionary with operation status and the user chat preference object or error details.
Supports both new polymorphic structure (focus_entity_type, focus_entity_id) and
legacy structure (focus_project_id, focus_workspace_id) for backward compatibility.
New parameters take precedence over legacy parameters.
"""
try:
stmt = select(UserChatPreference).where(UserChatPreference.user_id == user_id).where(UserChatPreference.chat_id == chat_id) # type: ignore[union-attr,arg-type]
result = await db.execute(stmt)
existing_user_chat_preference = result.scalar_one_or_none()
# Determine focus context: prioritize new polymorphic params, fall back to legacy
final_focus_entity_type = focus_entity_type
final_focus_entity_id = focus_entity_id
if not final_focus_entity_type and not final_focus_entity_id:
# Fall back to legacy parameters
if focus_project_id:
final_focus_entity_type = "project"
final_focus_entity_id = focus_project_id
elif focus_workspace_id:
final_focus_entity_type = "workspace"
final_focus_entity_id = focus_workspace_id
if existing_user_chat_preference:
if is_focus_enabled is not None:
existing_user_chat_preference.is_focus_enabled = is_focus_enabled
if final_focus_entity_type is not None:
existing_user_chat_preference.focus_entity_type = final_focus_entity_type
if final_focus_entity_id is not None:
existing_user_chat_preference.focus_entity_id = final_focus_entity_id
# Also update legacy fields for backward compatibility during migration
if focus_project_id is not None:
existing_user_chat_preference.focus_project_id = focus_project_id
if focus_workspace_id is not None:
existing_user_chat_preference.focus_workspace_id = focus_workspace_id
if mode is not None:
existing_user_chat_preference.mode = mode
db.add(existing_user_chat_preference)
await db.commit()
return {"message": "success", "user_chat_preference": existing_user_chat_preference}
@@ -1524,10 +1762,17 @@ async def upsert_user_chat_preference(
}
if is_focus_enabled is not None:
new_user_chat_preference_kwargs["is_focus_enabled"] = is_focus_enabled
if final_focus_entity_type is not None:
new_user_chat_preference_kwargs["focus_entity_type"] = final_focus_entity_type
if final_focus_entity_id is not None:
new_user_chat_preference_kwargs["focus_entity_id"] = final_focus_entity_id
# Also set legacy fields for backward compatibility during migration
if focus_project_id is not None:
new_user_chat_preference_kwargs["focus_project_id"] = focus_project_id
if focus_workspace_id is not None:
new_user_chat_preference_kwargs["focus_workspace_id"] = focus_workspace_id
if mode is not None:
new_user_chat_preference_kwargs["mode"] = mode
new_user_chat_preference = UserChatPreference(**new_user_chat_preference_kwargs)
db.add(new_user_chat_preference)
await db.commit()
@@ -1536,3 +1781,40 @@ async def upsert_user_chat_preference(
await db.rollback()
log.error(f"Error upserting user chat preference: {e}")
return {"message": "error", "error": str(e)}
async def get_project_from_chat_preference(chat_id: str, user_id: str, db: AsyncSession) -> Tuple[Optional[str], bool]:
# Get is_project_chat from chat record (with explicit commit/refresh to avoid race condition)
# Refresh the session to ensure we see any committed chat records
await db.commit()
stmt = select(Chat).where(Chat.id == chat_id) # type: ignore[arg-type]
result = await db.execute(stmt)
chat = result.scalar_one_or_none()
is_project_chat = chat.is_project_chat if chat else False
# Get project_id from UserChatPreference
project_id = None
try:
pref_stmt = select(UserChatPreference).where(
UserChatPreference.chat_id == chat_id, # type: ignore[arg-type]
UserChatPreference.user_id == user_id, # type: ignore[arg-type]
)
result = await db.execute(pref_stmt)
pref = result.scalar_one_or_none()
if pref:
# Check polymorphic focus first
if pref.focus_entity_type == FocusEntityType.PROJECT.value and pref.focus_entity_id:
project_id = pref.focus_entity_id
# Fallback to legacy field
elif pref.focus_project_id:
project_id = pref.focus_project_id
if project_id:
log.info(f"🔧 Found project_id in UserChatPreference: {project_id}")
except Exception as e:
log.error(f"Error fetching UserChatPreference: {e}")
project_id = str(project_id) if project_id else None
return project_id, is_project_chat
apps/pi/pi/services/retrievers/pg_store/message.py
+30 -9
View File
@@ -19,6 +19,7 @@ from pi.app.models import Message
from pi.app.models import MessageFeedback
from pi.app.models import MessageFlowStep
from pi.app.models import MessageMeta
from pi.app.models import UserChatPreference
from pi.app.models.enums import ExecutionStatus
from pi.app.models.enums import FlowStepType
from pi.app.models.enums import MessageFeedbackTypeChoices
@@ -200,6 +201,7 @@ async def upsert_message(
reasoning: Optional[str] = None,
parsed_content: Optional[str] = None,
workspace_slug: Optional[str] = None,
source: Optional[str] = None,
) -> Dict[str, Any]:
"""
Creates a new message or updates an existing one.
@@ -243,6 +245,8 @@ async def upsert_message(
existing_message.parsed_content = parsed_content
if workspace_slug is not None:
existing_message.workspace_slug = workspace_slug
if source is not None:
existing_message.source = source
# updated_at will be handled by SQLAlchemy
db.add(existing_message)
await db.commit()
@@ -273,6 +277,7 @@ async def upsert_message(
llm_model_id=None, # Add the missing required parameter
reasoning=reasoning or "",
workspace_slug=workspace_slug,
source=source,
)
db.add(new_message)
await db.commit()
@@ -494,13 +499,14 @@ async def upsert_message_meta(
workspace_slug = msg.workspace_slug
if meta:
# --- UPDATE ---
meta.input_text_tokens = in_tokens
meta.input_text_price = input_cost_usd
meta.output_text_tokens = out_tokens
meta.output_text_price = output_cost_usd
meta.cached_input_text_tokens = cached_in_tokens
meta.cached_input_text_price = cached_input_cost_usd
# --- UPDATE --- Aggregate tokens and costs instead of overwriting
# This handles cases where multiple LLM calls occur for the same message_id + step_type
meta.input_text_tokens = (meta.input_text_tokens or 0) + in_tokens
meta.input_text_price = (meta.input_text_price or 0.0) + input_cost_usd
meta.output_text_tokens = (meta.output_text_tokens or 0) + out_tokens
meta.output_text_price = (meta.output_text_price or 0.0) + output_cost_usd
meta.cached_input_text_tokens = (meta.cached_input_text_tokens or 0) + cached_in_tokens
meta.cached_input_text_price = (meta.cached_input_text_price or 0.0) + cached_input_cost_usd
meta.llm_model_id = llm_model_id
meta.llm_model_pricing_id = pricing_id
if workspace_slug is not None:
@@ -636,6 +642,20 @@ async def reconstruct_chat_request_from_message(db: AsyncSession, user_message:
attachments = attachments_result.scalars().all()
attachment_ids = [att.attachment_id for att in attachments] if attachments else []
# Get user chat preferences to retrieve mode
preferences_stmt = (
select(UserChatPreference)
.where(UserChatPreference.chat_id == user_message.chat_id) # type: ignore[arg-type]
.where(UserChatPreference.deleted_at.is_(None)) # type: ignore[union-attr]
)
preferences_result = await db.execute(preferences_stmt)
user_preferences = preferences_result.scalar_one_or_none()
mode = user_preferences.mode if user_preferences and user_preferences.mode else "ask"
project_id = user_preferences.focus_project_id if user_preferences and user_preferences.focus_project_id else None
# Get source from message, default to "web" if not set
source = user_message.source or "web"
# Reconstruct ChatRequest
return ChatRequest(
query=user_message.content or "",
@@ -648,8 +668,9 @@ async def reconstruct_chat_request_from_message(db: AsyncSession, user_message:
workspace_slug=chat.workspace_slug if chat else None,
workspace_in_context=chat.workspace_in_context if chat else False,
is_project_chat=chat.is_project_chat if chat else False,
project_id=None, # Will be resolved from context if needed
project_id=project_id,
attachment_ids=attachment_ids,
context={"token_id": str(user_message.id)}, # Reuse same message ID
source="web",
source=source,
mode=mode,
)
apps/pi/pi/services/schemas/chat.py
+37 -16
View File
@@ -10,24 +10,35 @@ from pydantic import BaseModel
from pydantic import Field
class Agents(str, Enum):
GENERIC_AGENT = "generic_agent"
PLANE_STRUCTURED_DATABASE_AGENT = "plane_structured_database_agent"
PLANE_VECTOR_SEARCH_AGENT = "plane_vector_search_agent"
PLANE_PAGES_AGENT = "plane_pages_agent"
PLANE_DOCS_AGENT = "plane_docs_agent"
PLANE_ACTION_EXECUTOR_AGENT = "plane_action_executor_agent"
class RetrievalTools(str, Enum):
"""Enumeration of available retrieval tools (formerly called agents)."""
STRUCTURED_DB_TOOL = "structured_db_tool"
VECTOR_SEARCH_TOOL = "vector_search_tool"
PAGES_SEARCH_TOOL = "pages_search_tool"
DOCS_SEARCH_TOOL = "docs_search_tool"
ACTION_EXECUTOR_TOOL = "action_executor_tool"
class AgentQuery(BaseModel):
agent: Agents
query: str = Field(..., description="The decomposed query for this specific agent")
# Legacy alias for backward compatibility during migration
Agents = RetrievalTools
class ToolQuery(BaseModel):
"""A query targeted at a specific retrieval tool."""
tool: RetrievalTools
query: str = Field(..., description="The decomposed query for this specific tool")
# Legacy alias for backward compatibility during migration
AgentQuery = ToolQuery
class RouteQuery(BaseModel):
"""Route a user query to the most relevant customer service agent(s) with decomposed queries."""
"""Route a user query to the most relevant retrieval tool(s) with decomposed queries."""
decomposed_queries: list[AgentQuery] = Field(..., description="List of agents with their corresponding decomposed queries")
decomposed_queries: list[ToolQuery] = Field(..., description="List of tools with their corresponding decomposed queries")
class RoutingResult(TypedDict):
@@ -36,13 +47,20 @@ class RoutingResult(TypedDict):
parsing_error: Exception | None
AgentQueryList: TypeAlias = list[AgentQuery]
ToolQueryList: TypeAlias = list[ToolQuery]
# Legacy alias for backward compatibility during migration
AgentQueryList = ToolQueryList
class AgentOrder(BaseModel):
"""Provide the order in which the selected data retrieval agents should be executed."""
class ToolOrder(BaseModel):
"""Provide the order in which the selected data retrieval tools should be executed."""
ordered_agents: List[str] = Field(..., description="List of ordered agent names")
ordered_tools: List[str] = Field(..., description="List of ordered tool names")
# Legacy alias for backward compatibility during migration
AgentOrder = ToolOrder
class QueryFlowStore(TypedDict):
@@ -61,6 +79,8 @@ class QueryFlowStore(TypedDict):
workspace_in_context: bool
project_id: str
workspace_id: str
# Optional: current step order to align sub-step persistence without DB lookups
step_order: int
# --- Action Category Routing (for hierarchical actions) ---
@@ -88,6 +108,7 @@ class ActionCategorySelection(BaseModel):
"properties",
"types",
"worklogs",
"retrieval_tools",
]
rationale: Optional[str] = None
apps/pi/pi/tests/example_slack_answer_request.py
+145
View File
@@ -0,0 +1,145 @@
"""
Example Python requests code to hit the /slack/answer/ endpoint.
This endpoint is located at: {base_path}/api/v1/chat/slack/answer/
"""
import json
from uuid import UUID
from uuid import uuid4
import requests
# Configuration
BASE_URL = "http://localhost:8001" # Update with your actual base URL
BASE_PATH = "" # Update if your API has a base path (e.g., "/pi")
AUTHORIZATION_TOKEN = "plane_api_c9a0aae9011c483bae231b655afac38d" # Your Plane access token or API key
# AUTHORIZATION_TOKEN = "plane_api_c9a0aae9011c483bae231b655afac38d" # Your Plane access token or API key
# Endpoint URL
endpoint_url = f"{BASE_URL}{BASE_PATH}/api/v1/chat/slack/answer/"
llm = "gpt-4.1"
# llm = 'gpt-5-standard'
# llm = 'gpt-5-fast'
# llm = 'claude-sonnet-4'
user_id = UUID("c1cc72be-3b92-445d-b223-6bcd94876925")
# user_id = UUID("2225dd79-dffe-4272-9748-3f88548daddb")
# user_id = UUID("6baf27f5-1bfd-4acb-aa2f-333cd23130e0")
workspace_id = "42200c3d-1225-4cb5-88ba-c88af80430f2"
# workspace_id = "cd4ab5a2-1a5f-4516-a6c6-8da1a9fa5be4"
workspace_slug = None
# workspace_id = "23d3924f-db82-4377-b8d3-6ff643af6b9a"
# workspace_id = "31f16e25-535e-404f-946e-3ffbe52695a2"
# workspace_id = "1559196e-fdd9-4651-8a78-33e30cbf5f72"
# workspace_id = None
is_temp = False
workspace_in_context = True
is_new = True
pi_sidebar_open = True
sidebar_open_url = r"piworkspace/projects/3719d6d4-42d6-4bf5-bfe4-85dc33e07958/cycles/28ec0918-9095-42c7-80ae-864816d8dbd2/"
is_project_chat = False
# project_id = "40bc0eaf-2d34-457a-ae5b-d4040c42edcc"
project_id = "47ed2e28-3e91-4ba8-a463-4a117941e2d4"
# project_id=None
# project_id = "743525c9-7ac3-4570-8262-e63330175ed4"
context = {"first_name": "sunder", "last_name": "Chintada"}
query = """
This was the previous conversation history: \n\n
===================\n\n
User: give me one issue assigned to me in PALL project
Assistant: There are currently no work-items assigned to you in the PALL project. If you need help finding work-items in another project or have a different request, just let me know!
Assistant: There are currently no work-items assigned to you in the PALL project. If you need help finding work-items in another project or want to adjust your search, just let me know!
User: create a task in PALL project title - "Hello task" and assign it to me
Assistant: Successfully created work item 'Hello task' - <http://127.0.0.1:3000/integrations03/browse/PALL-4/>
User: do I have any task assigned to me in PALL project. If yes, show it
Assistant: Yes, you currently have one task assigned to you in the PALL project:
• *Task Name:* Hello task
• *Unique Key:* Not provided in the data
• *Priority:* None
• *State:* Backlog
If you need more details or want to take action on this task, just let me know!
User: what project are we discussing in this thread
Assistant: It seems like you're referring to a previous conversation about the PALL project. If you have any specific questions or need further information about the PALL project, feel free to ask!
User: create a task with title Test World
Assistant: Successfully created work item 'Test World' - <http://127.0.0.1:3000/integrations03/browse/PALL-5/>
User: what are we discussing?
Assistant: I'm sorry, but I don't have access to past conversation history. However, I can help with any questions or information you need right now. Feel free to ask!
User: try again
Assistant: I'm here to help with any questions or tasks you have. However, I don't have access to past conversation history. If you have any specific questions or need assistance, feel free to let me know!
User: past conversation is available under `This was the previous conversation history` what are we discussing
Assistant: I'm sorry, but I don't have access to the previous conversation history. However, if you provide me with some details or context about the project and work items, I'd be happy to help you summarize or discuss them further!
User: summarize this thread
Assistant: I'm sorry, but I can't provide a summary of the conversation history as there is no prior conversation content available. If you have specific questions or need assistance with a project, feel free to let me know!
===================\n\n
πCurrent user messageπ:\
What are we discussing in this thread?
""" # noqa E501
query = """
πCurrent user messageπ:
create a task with title Test World in Starlette project
"""
# Request payload - ChatRequest schema
payload = {
"query": query,
"llm": llm, # Optional: defaults to DEFAULT_LLM if not provided
"is_new": True,
"user_id": str(user_id), # Optional: UUID as string
"chat_id": str(uuid4()), # Optional: UUID as string
"is_temp": False,
"workspace_in_context": True,
"workspace_id": str(workspace_id), # Optional: UUID as string or empty string
"project_id": None, # Optional: UUID as string or None
"context": {}, # Required: dictionary with any context data
"is_project_chat": False, # Optional
"pi_sidebar_open": False, # Optional
"workspace_slug": "switchertest", # Optional: workspace slug string
"attachment_ids": [], # Optional: list of UUID strings
"sidebar_open_url": "", # Optional
"source": "app", # Optional: source identifier
}
# Headers
headers = {
"Authorization": f"Bearer {AUTHORIZATION_TOKEN}", # Can also be just the token without "Bearer "
"Content-Type": "application/json",
}
# Make the request
try:
response = requests.post(
endpoint_url,
json=payload,
headers=headers,
timeout=300, # Adjust timeout as needed (streaming can take time)
)
# Check response status
response.raise_for_status()
# Parse JSON response
result = response.json()
print("Response Status:", response.status_code)
print("\nResponse Body:")
print(json.dumps(result, indent=2))
# Access response fields
print("\n--- Response Fields ---")
print(f"Response Type: {result.get("response_type")}")
print(f"Response: {result.get("response")}")
print(f"Actions Data: {result.get("actions_data")}")
print(f"Context: {result.get("context")}")
print(f"Clarification Data: {result.get("clarification_data")}")
except requests.exceptions.HTTPError as e:
print(f"HTTP Error: {e}")
print(f"Response: {e.response.text}")
except requests.exceptions.RequestException as e:
print(f"Request Error: {e}")
apps/pi/pi/tests/services/test_cycles_add_work_items.py
+74
View File
@@ -0,0 +1,74 @@
"""
Test for cycles.add_work_items method to verify the fix for the 'data' keyword argument error.
This test can be run with a real API token to verify the fix works correctly.
"""
import os
import pytest
from pi.services.actions.plane_sdk_adapter import PlaneSDKAdapter
class TestCyclesAddWorkItems:
"""Test cycles.add_work_items method."""
@pytest.mark.skipif(
not os.getenv("PLANE_API_TOKEN"),
reason="PLANE_API_TOKEN environment variable not set. Set it to run this test with real API.",
)
def test_add_work_items_to_cycle(self):
"""Test adding work items to a cycle with real API."""
# Get API token from environment
api_token = os.getenv("PLANE_API_TOKEN")
workspace_slug = os.getenv("PLANE_WORKSPACE_SLUG", "test-workspace")
project_id = os.getenv("PLANE_PROJECT_ID", "test-project")
cycle_id = os.getenv("PLANE_CYCLE_ID", "test-cycle")
work_item_id = os.getenv("PLANE_WORK_ITEM_ID", "test-work-item")
# Create adapter
adapter = PlaneSDKAdapter(access_token=api_token)
# Test adding work items to cycle
# This should work without the 'data' keyword argument error
result = adapter.add_cycle_work_items(
workspace_slug=workspace_slug,
project_id=project_id,
cycle_id=cycle_id,
issues=[work_item_id],
)
# Verify result
assert result is not None
print("✅ Successfully added work items to cycle!")
print(f" Result: {result}")
@pytest.mark.skipif(
not os.getenv("PLANE_API_TOKEN"),
reason="PLANE_API_TOKEN environment variable not set. Set it to run this test with real API.",
)
def test_add_work_items_to_cycle_with_multiple_issues(self):
"""Test adding multiple work items to a cycle."""
api_token = os.getenv("PLANE_API_TOKEN")
workspace_slug = os.getenv("PLANE_WORKSPACE_SLUG", "test-workspace")
project_id = os.getenv("PLANE_PROJECT_ID", "test-project")
cycle_id = os.getenv("PLANE_CYCLE_ID", "test-cycle")
work_item_ids = os.getenv("PLANE_WORK_ITEM_IDS", "item1,item2").split(",")
adapter = PlaneSDKAdapter(access_token=api_token)
result = adapter.add_cycle_work_items(
workspace_slug=workspace_slug,
project_id=project_id,
cycle_id=cycle_id,
issues=work_item_ids,
)
assert result is not None
print(f"✅ Successfully added {len(work_item_ids)} work items to cycle!")
print(f" Result: {result}")
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
apps/pi/pi/tests/services/test_plane_api_methods.py
-796
View File
@@ -1,796 +0,0 @@
"""
Comprehensive Test Suite for Plane API Methods Signature Validation.
This test suite validates that ALL 81 implemented Plane API methods have signatures
that exactly match the Plane SDK specifications. It covers 17 API categories including
the original 8 categories (36 methods) plus 9 new categories (45 methods).
Coverage:
- Projects API (4 methods)
- Work Items API (4 methods)
- Cycles API (10 methods)
- Labels API (2 methods)
- States API (2 methods)
- Modules API (8 methods)
- Assets API (5 methods)
- Users API (1 method)
- Intake API (5 methods) - NEW
- Members API (2 methods) - NEW
- Activity API (2 methods) - NEW
- Attachments API (4 methods) - NEW
- Comments API (5 methods) - NEW
- Links API (5 methods) - NEW
- Properties API (12 methods) - NEW
- Types API (5 methods) - NEW
- Worklogs API (5 methods) - NEW
This test suite does NOT make actual API calls but focuses on parameter validation,
method signature compliance, and SDK request object usage.
"""
from unittest.mock import AsyncMock
import pytest
from pi.services.actions.method_executor import MethodExecutor
# Configure pytest for async testing
pytest_plugins = ("pytest_asyncio",)
class TestPlaneAPIMethods:
"""Test class for validating Plane API method signatures against SDK specifications."""
@pytest.fixture
def method_executor(self):
"""Create a MethodExecutor instance with mocked execute method."""
# Create a mock MethodExecutor that returns success for all execute calls
mock_executor = AsyncMock()
mock_executor.execute = AsyncMock(return_value={"success": True, "data": "test_data"})
mock_executor.get_category_methods = MethodExecutor.get_category_methods
return mock_executor
@pytest.mark.asyncio
async def test_projects_api_method_signatures(self, method_executor):
"""Test Projects API method signatures match SDK specifications."""
# Test projects retrieve - SDK: retrieve_project(pk, slug)
result = await method_executor.execute("projects", "retrieve", pk="project-123", slug="test-workspace")
assert result["success"] is True
# Test projects update - SDK: update_project(pk, slug, patched_project_update_request=None)
update_data = {"name": "Updated Project", "description": "Updated description"}
result = await method_executor.execute(
"projects", "update", pk="project-123", slug="test-workspace", patched_project_update_request=update_data
)
assert result["success"] is True
# Test projects archive - SDK: archive_project(project_id, slug)
result = await method_executor.execute("projects", "archive", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test projects unarchive - SDK: unarchive_project(project_id, slug)
result = await method_executor.execute("projects", "unarchive", project_id="project-123", slug="test-workspace")
assert result["success"] is True
@pytest.mark.asyncio
async def test_workitems_api_method_signatures(self, method_executor):
"""Test Work Items API method signatures match SDK specifications."""
# Test workitems list with all optional parameters
result = await method_executor.execute(
"workitems",
"list",
project_id="project-123",
slug="test-workspace",
cursor="cursor-123",
expand="assignees,labels",
external_id="ext-123",
external_source="jira",
fields="name,description",
order_by="-created_at",
per_page=50,
)
assert result["success"] is True
# Test workitems retrieve with all optional parameters
result = await method_executor.execute(
"workitems",
"retrieve",
pk="issue-123",
project_id="project-123",
slug="test-workspace",
expand="assignees,labels",
external_id="ext-123",
external_source="jira",
fields="name,description",
order_by="-created_at",
)
assert result["success"] is True
# Test workitems search
result = await method_executor.execute(
"workitems", "search", search="bug report", slug="test-workspace", limit=25, project_id="project-123", workspace_search=True
)
assert result["success"] is True
# Test workitems get_workspace
result = await method_executor.execute("workitems", "get_workspace", issue_identifier="123", project_identifier="TEST", slug="test-workspace")
assert result["success"] is True
@pytest.mark.asyncio
async def test_cycles_api_method_signatures(self, method_executor):
"""Test Cycles API method signatures match SDK specifications."""
# Test cycles retrieve
result = await method_executor.execute("cycles", "retrieve", pk="cycle-123", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test cycles update with request object
update_data = {"name": "Updated Sprint", "description": "Updated description"}
result = await method_executor.execute(
"cycles", "update", pk="cycle-123", project_id="project-123", slug="test-workspace", patched_cycle_update_request=update_data
)
assert result["success"] is True
# Test cycles add_work_items with request object
request_data = {"issues": ["issue-1", "issue-2"]}
result = await method_executor.execute(
"cycles",
"add_work_items",
cycle_id="cycle-123",
project_id="project-123",
slug="test-workspace",
cycle_issue_request_request=request_data,
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_labels_api_method_signatures(self, method_executor):
"""Test Labels API method signatures match SDK specifications."""
# Test labels retrieve
result = await method_executor.execute("labels", "retrieve", pk="label-123", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test labels update with request object
update_data = {"name": "Updated Bug", "color": "#ff0000", "description": "Updated description"}
result = await method_executor.execute(
"labels", "update", pk="label-123", project_id="project-123", slug="test-workspace", patched_label_create_update_request=update_data
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_states_api_method_signatures(self, method_executor):
"""Test States API method signatures match SDK specifications."""
# Test states retrieve - note the parameter order: project_id, slug, state_id
result = await method_executor.execute("states", "retrieve", project_id="project-123", slug="test-workspace", state_id="state-123")
assert result["success"] is True
# Test states update with request object
update_data = {"name": "In Review", "color": "#ffaa00", "description": "Work items under review"}
result = await method_executor.execute(
"states", "update", project_id="project-123", slug="test-workspace", state_id="state-123", patched_state_request=update_data
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_modules_api_method_signatures(self, method_executor):
"""Test Modules API method signatures match SDK specifications."""
# Test modules retrieve
result = await method_executor.execute("modules", "retrieve", pk="module-123", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test modules update with request object
update_data = {"name": "Updated Module", "description": "Updated description"}
result = await method_executor.execute(
"modules", "update", pk="module-123", project_id="project-123", slug="test-workspace", patched_module_update_request=update_data
)
assert result["success"] is True
# Test modules add_work_items with request object
request_data = {"issues": ["issue-1", "issue-2"]}
result = await method_executor.execute(
"modules",
"add_work_items",
module_id="module-123",
project_id="project-123",
slug="test-workspace",
module_issue_request_request=request_data,
)
assert result["success"] is True
# Test modules remove_work_item - note parameter order: issue_id, module_id, project_id, slug
result = await method_executor.execute(
"modules", "remove_work_item", issue_id="issue-123", module_id="module-123", project_id="project-123", slug="test-workspace"
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_assets_api_method_signatures(self, method_executor):
"""Test Assets API method signatures match SDK specifications."""
# Test assets create_user_upload with request object
request_data = {"name": "test_image.png", "size": 1024, "type": "image"}
result = await method_executor.execute("assets", "create_user_upload", user_asset_upload_request=request_data)
assert result["success"] is True
# Test assets get_generic
result = await method_executor.execute("assets", "get_generic", asset_id="asset-123", slug="test-workspace")
assert result["success"] is True
# Test assets update_generic with request object
update_data = {"name": "Updated Asset", "description": "Updated description"}
result = await method_executor.execute(
"assets", "update_generic", asset_id="asset-123", slug="test-workspace", patched_generic_asset_update_request=update_data
)
assert result["success"] is True
# Test assets update_user with request object
update_data = {"name": "Updated User Asset"}
result = await method_executor.execute("assets", "update_user", asset_id="asset-123", patched_asset_update_request=update_data)
assert result["success"] is True
# Test assets delete_user
result = await method_executor.execute("assets", "delete_user", asset_id="asset-123")
assert result["success"] is True
@pytest.mark.asyncio
async def test_users_api_method_signatures(self, method_executor):
"""Test Users API method signatures match SDK specifications."""
# Test users get_current - no parameters required
result = await method_executor.execute("users", "get_current")
assert result["success"] is True
@pytest.mark.asyncio
async def test_intake_api_method_signatures(self, method_executor):
"""Test Intake API method signatures match SDK specifications."""
# Test intake create with request object
intake_request = {"name": "New Intake Item", "description_html": "<p>Test description</p>", "priority": "medium"}
result = await method_executor.execute(
"intake", "create", project_id="project-123", slug="test-workspace", intake_issue_create_request=intake_request
)
assert result["success"] is True
# Test intake delete
result = await method_executor.execute("intake", "delete", intake_id="intake-123", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test intake list with pagination
result = await method_executor.execute("intake", "list", project_id="project-123", slug="test-workspace", cursor="cursor-123", per_page=20)
assert result["success"] is True
# Test intake retrieve
result = await method_executor.execute("intake", "retrieve", intake_id="intake-123", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test intake update with request object
update_data = {"name": "Updated Intake", "priority": "high"}
result = await method_executor.execute(
"intake",
"update",
intake_id="intake-123",
project_id="project-123",
slug="test-workspace",
patched_intake_issue_update_request=update_data,
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_members_api_method_signatures(self, method_executor):
"""Test Members API method signatures match SDK specifications."""
# Test get_project_members
result = await method_executor.execute("members", "get_project_members", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test get_workspace_members
result = await method_executor.execute("members", "get_workspace_members", slug="test-workspace")
assert result["success"] is True
@pytest.mark.asyncio
async def test_activity_api_method_signatures(self, method_executor):
"""Test Work Item Activity API method signatures match SDK specifications."""
# Test list_work_item_activities with all optional parameters
result = await method_executor.execute(
"activity",
"list",
issue_id="issue-123",
project_id="project-123",
slug="test-workspace",
cursor="cursor-123",
expand="actor,issue",
fields="activity_type,created_at",
order_by="-created_at",
per_page=20,
)
assert result["success"] is True
# Test retrieve_work_item_activity
result = await method_executor.execute(
"activity", "retrieve", activity_id="activity-123", issue_id="issue-123", project_id="project-123", slug="test-workspace"
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_attachments_api_method_signatures(self, method_executor):
"""Test Work Item Attachments API method signatures match SDK specifications."""
# Test create_work_item_attachment with request object
attachment_request = {"name": "test_file.pdf", "size": 2048, "asset_type": "attachment"}
result = await method_executor.execute(
"attachments",
"create",
issue_id="issue-123",
project_id="project-123",
slug="test-workspace",
issue_attachment_upload_request=attachment_request,
)
assert result["success"] is True
# Test delete_work_item_attachment
result = await method_executor.execute(
"attachments", "delete", attachment_id="attachment-123", issue_id="issue-123", project_id="project-123", slug="test-workspace"
)
assert result["success"] is True
# Test list_work_item_attachments
result = await method_executor.execute("attachments", "list", issue_id="issue-123", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test retrieve_work_item_attachment
result = await method_executor.execute(
"attachments", "retrieve", attachment_id="attachment-123", issue_id="issue-123", project_id="project-123", slug="test-workspace"
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_comments_api_method_signatures(self, method_executor):
"""Test Work Item Comments API method signatures match SDK specifications."""
# Test create_work_item_comment with request object
comment_request = {"comment_html": "<p>This is a test comment</p>"}
result = await method_executor.execute(
"comments", "create", issue_id="issue-123", project_id="project-123", slug="test-workspace", issue_comment_create_request=comment_request
)
assert result["success"] is True
# Test delete_work_item_comment
result = await method_executor.execute(
"comments", "delete", comment_id="comment-123", issue_id="issue-123", project_id="project-123", slug="test-workspace"
)
assert result["success"] is True
# Test list_work_item_comments with pagination
result = await method_executor.execute(
"comments",
"list",
issue_id="issue-123",
project_id="project-123",
slug="test-workspace",
cursor="cursor-123",
expand="actor",
fields="comment_html,created_at",
order_by="-created_at",
per_page=20,
)
assert result["success"] is True
# Test retrieve_work_item_comment
result = await method_executor.execute(
"comments", "retrieve", comment_id="comment-123", issue_id="issue-123", project_id="project-123", slug="test-workspace"
)
assert result["success"] is True
# Test update_work_item_comment with request object
update_data = {"comment_html": "<p>Updated comment</p>"}
result = await method_executor.execute(
"comments",
"update",
comment_id="comment-123",
issue_id="issue-123",
project_id="project-123",
slug="test-workspace",
patched_issue_comment_create_request=update_data,
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_links_api_method_signatures(self, method_executor):
"""Test Work Item Links API method signatures match SDK specifications."""
# Test create_work_item_link with request object
link_request = {"url": "https://example.com/reference", "title": "External Reference"}
result = await method_executor.execute(
"links", "create", issue_id="issue-123", project_id="project-123", slug="test-workspace", issue_link_create_request=link_request
)
assert result["success"] is True
# Test delete_work_item_link
result = await method_executor.execute(
"links", "delete", link_id="link-123", issue_id="issue-123", project_id="project-123", slug="test-workspace"
)
assert result["success"] is True
# Test list_work_item_links with pagination
result = await method_executor.execute(
"links",
"list",
issue_id="issue-123",
project_id="project-123",
slug="test-workspace",
cursor="cursor-123",
expand="metadata",
fields="url,title",
order_by="-created_at",
per_page=20,
)
assert result["success"] is True
# Test retrieve_work_item_link
result = await method_executor.execute(
"links",
"retrieve",
link_id="link-123",
issue_id="issue-123",
project_id="project-123",
slug="test-workspace",
expand="metadata",
fields="url,title",
)
assert result["success"] is True
# Test update_issue_link with request object
update_data = {"url": "https://updated-example.com", "title": "Updated Reference"}
result = await method_executor.execute(
"links",
"update",
link_id="link-123",
issue_id="issue-123",
project_id="project-123",
slug="test-workspace",
patched_issue_link_update_request=update_data,
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_properties_api_method_signatures(self, method_executor):
"""Test Work Item Properties API method signatures match SDK specifications."""
# Test create_issue_property with request object
property_request = {"name": "Priority Level", "description": "Custom priority field"}
result = await method_executor.execute(
"properties", "create", project_id="project-123", slug="test-workspace", type_id="type-123", issue_property_api_request=property_request
)
assert result["success"] is True
# Test create_issue_property_option with request object
option_request = {"name": "High Priority", "value": "high"}
result = await method_executor.execute(
"properties",
"create_option",
project_id="project-123",
property_id="property-123",
slug="test-workspace",
type_id="type-123",
issue_property_option_api_request=option_request,
)
assert result["success"] is True
# Test create_issue_property_value with request object
value_request = {"value": "high"}
result = await method_executor.execute(
"properties",
"create_value",
issue_id="issue-123",
project_id="project-123",
property_id="property-123",
slug="test-workspace",
type_id="type-123",
issue_property_value_api_request=value_request,
)
assert result["success"] is True
# Test list_issue_properties
result = await method_executor.execute("properties", "list", project_id="project-123", slug="test-workspace", type_id="type-123")
assert result["success"] is True
# Test retrieve_issue_property
result = await method_executor.execute(
"properties", "retrieve", project_id="project-123", property_id="property-123", slug="test-workspace", type_id="type-123"
)
assert result["success"] is True
# Test update_issue_property with request object
update_data = {"name": "Updated Priority", "description": "Updated description"}
result = await method_executor.execute(
"properties",
"update",
project_id="project-123",
property_id="property-123",
slug="test-workspace",
type_id="type-123",
patched_issue_property_api_request=update_data,
)
assert result["success"] is True
# Test delete_issue_property
result = await method_executor.execute(
"properties", "delete", project_id="project-123", property_id="property-123", slug="test-workspace", type_id="type-123"
)
assert result["success"] is True
# Test delete_issue_property_option
result = await method_executor.execute(
"properties",
"delete_option",
option_id="option-123",
project_id="project-123",
property_id="property-123",
slug="test-workspace",
type_id="type-123",
)
assert result["success"] is True
# Test list_issue_property_options
result = await method_executor.execute(
"properties", "list_options", project_id="project-123", property_id="property-123", slug="test-workspace", type_id="type-123"
)
assert result["success"] is True
# Test list_issue_property_values
result = await method_executor.execute(
"properties", "list_values", issue_id="issue-123", project_id="project-123", slug="test-workspace", type_id="type-123"
)
assert result["success"] is True
# Test retrieve_issue_property_option
result = await method_executor.execute(
"properties",
"retrieve_option",
option_id="option-123",
project_id="project-123",
property_id="property-123",
slug="test-workspace",
type_id="type-123",
)
assert result["success"] is True
# Test update_issue_property_option with request object
update_data = {"name": "Critical Priority", "value": "critical"}
result = await method_executor.execute(
"properties",
"update_option",
option_id="option-123",
project_id="project-123",
property_id="property-123",
slug="test-workspace",
type_id="type-123",
patched_issue_property_option_api_request=update_data,
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_types_api_method_signatures(self, method_executor):
"""Test Work Item Types API method signatures match SDK specifications."""
# Test create_issue_type with request object
type_request = {"name": "Epic", "description": "Large feature or initiative"}
result = await method_executor.execute(
"types", "create", project_id="project-123", slug="test-workspace", issue_type_api_request=type_request
)
assert result["success"] is True
# Test delete_issue_type
result = await method_executor.execute("types", "delete", project_id="project-123", slug="test-workspace", type_id="type-123")
assert result["success"] is True
# Test list_issue_types
result = await method_executor.execute("types", "list", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test retrieve_issue_type
result = await method_executor.execute("types", "retrieve", project_id="project-123", slug="test-workspace", type_id="type-123")
assert result["success"] is True
# Test update_issue_type with request object
update_data = {"name": "User Story", "description": "Updated description"}
result = await method_executor.execute(
"types", "update", project_id="project-123", slug="test-workspace", type_id="type-123", patched_issue_type_api_request=update_data
)
assert result["success"] is True
@pytest.mark.asyncio
async def test_worklogs_api_method_signatures(self, method_executor):
"""Test Work Item Worklogs API method signatures match SDK specifications."""
# Test create_issue_worklog with request object
worklog_request = {
"description": "Worked on feature implementation",
"duration": 120, # 2 hours in minutes
}
result = await method_executor.execute(
"worklogs", "create", issue_id="issue-123", project_id="project-123", slug="test-workspace", issue_work_log_api_request=worklog_request
)
assert result["success"] is True
# Test delete_issue_worklog
result = await method_executor.execute("worklogs", "delete", project_id="project-123", slug="test-workspace", worklog_id="worklog-123")
assert result["success"] is True
# Test get_project_worklog_summary
result = await method_executor.execute("worklogs", "get_summary", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test list_issue_worklogs
result = await method_executor.execute("worklogs", "list", issue_id="issue-123", project_id="project-123", slug="test-workspace")
assert result["success"] is True
# Test update_issue_worklog with request object
update_data = {
"description": "Updated work description",
"duration": 180, # 3 hours in minutes
}
result = await method_executor.execute(
"worklogs",
"update",
project_id="project-123",
slug="test-workspace",
worklog_id="worklog-123",
patched_issue_work_log_api_request=update_data,
)
assert result["success"] is True
def test_all_81_methods_available(self, method_executor):
"""Verify that all 81 methods are available through method_executor."""
# Get available methods for each category (including all new categories)
categories = [
"projects",
"workitems",
"cycles",
"labels",
"states",
"modules",
"assets",
"users",
"intake",
"members",
"activity",
"attachments",
"comments",
"links",
"properties",
"types",
"worklogs",
]
total_methods = 0
for category in categories:
methods = method_executor.get_category_methods(category)
total_methods += len(methods)
assert len(methods) > 0, f"No methods found for category {category}"
# We implemented 81 methods total (36 original + 45 new), should have at least that many
assert total_methods >= 81, f"Expected at least 81 methods, got {total_methods}"
def test_sdk_compliance_documentation(self):
"""Document that all implemented methods are SDK compliant."""
# This test documents our SDK compliance achievement
sdk_compliant_methods = [
# Projects API (4 methods)
"projects.retrieve",
"projects.update",
"projects.archive",
"projects.unarchive",
# Work Items API (4 methods)
"workitems.list",
"workitems.retrieve",
"workitems.search",
"workitems.get_workspace",
# Cycles API (10 methods)
"cycles.retrieve",
"cycles.update",
"cycles.archive",
"cycles.unarchive",
"cycles.list_archived",
"cycles.add_work_items",
"cycles.list_work_items",
"cycles.retrieve_work_item",
"cycles.remove_work_item",
"cycles.transfer_work_items",
# Labels API (2 methods)
"labels.retrieve",
"labels.update",
# States API (2 methods)
"states.retrieve",
"states.update",
# Modules API (8 methods)
"modules.retrieve",
"modules.update",
"modules.archive",
"modules.unarchive",
"modules.list_archived",
"modules.add_work_items",
"modules.list_work_items",
"modules.remove_work_item",
# Assets API (5 methods)
"assets.create_user_upload",
"assets.get_generic",
"assets.update_generic",
"assets.update_user",
"assets.delete_user",
# Users API (1 method)
"users.get_current",
# Intake API (5 methods) - NEW
"intake.create",
"intake.delete",
"intake.list",
"intake.retrieve",
"intake.update",
# Members API (2 methods) - NEW
"members.get_project_members",
"members.get_workspace_members",
# Activity API (2 methods) - NEW
"activity.list",
"activity.retrieve",
# Attachments API (4 methods) - NEW
"attachments.create",
"attachments.delete",
"attachments.list",
"attachments.retrieve",
# Comments API (5 methods) - NEW
"comments.create",
"comments.delete",
"comments.list",
"comments.retrieve",
"comments.update",
# Links API (5 methods) - NEW
"links.create",
"links.delete",
"links.list",
"links.retrieve",
"links.update",
# Properties API (12 methods) - NEW
"properties.create",
"properties.create_option",
"properties.create_value",
"properties.delete",
"properties.delete_option",
"properties.list",
"properties.list_options",
"properties.list_values",
"properties.retrieve",
"properties.retrieve_option",
"properties.update",
"properties.update_option",
# Types API (5 methods) - NEW
"types.create",
"types.delete",
"types.list",
"types.retrieve",
"types.update",
# Worklogs API (5 methods) - NEW
"worklogs.create",
"worklogs.delete",
"worklogs.get_summary",
"worklogs.list",
"worklogs.update",
]
# Verify we have implemented 81 methods (36 original + 45 new)
assert len(sdk_compliant_methods) == 81, f"Expected 81 SDK compliant methods, documented {len(sdk_compliant_methods)}"
print("✅ All 81 Plane API methods are SDK compliant and tested!")
if __name__ == "__main__":
pytest.main([__file__, "-v"])
apps/pi/pi/tests/services/test_pydantic_compatibility.py
-308
View File
@@ -1,308 +0,0 @@
"""
Test Suite for Pydantic v1/v2 Compatibility.
This test suite specifically validates that the pydantic compatibility fixes
in PlaneSDKAdapter work correctly by testing request object creation.
"""
from unittest.mock import AsyncMock
import pytest
from pi.services.actions.plane_sdk_adapter import PlaneSDKAdapter
# Configure pytest for async testing
pytest_plugins = ("pytest_asyncio",)
class TestPydanticCompatibility:
"""Test class for validating pydantic v1/v2 compatibility."""
@pytest.fixture
def mock_sdk_apis(self):
"""Create mock SDK API objects."""
mock_intake = AsyncMock()
mock_intake.create_intake_work_item = AsyncMock(return_value={"id": "intake-123"})
mock_intake.update_intake_work_item = AsyncMock(return_value={"id": "intake-123"})
mock_attachments = AsyncMock()
mock_attachments.create_work_item_attachment = AsyncMock(return_value={"id": "attachment-123"})
mock_comments = AsyncMock()
mock_comments.create_work_item_comment = AsyncMock(return_value={"id": "comment-123"})
mock_comments.update_work_item_comment = AsyncMock(return_value={"id": "comment-123"})
mock_links = AsyncMock()
mock_links.create_work_item_link = AsyncMock(return_value={"id": "link-123"})
mock_links.update_issue_link = AsyncMock(return_value={"id": "link-123"})
mock_properties = AsyncMock()
mock_properties.create_issue_property = AsyncMock(return_value={"id": "property-123"})
mock_properties.update_issue_property = AsyncMock(return_value={"id": "property-123"})
mock_types = AsyncMock()
mock_types.create_issue_type = AsyncMock(return_value={"id": "type-123"})
mock_types.update_issue_type = AsyncMock(return_value={"id": "type-123"})
mock_worklogs = AsyncMock()
mock_worklogs.create_issue_worklog = AsyncMock(return_value={"id": "worklog-123"})
mock_worklogs.update_issue_worklog = AsyncMock(return_value={"id": "worklog-123"})
return {
"intake": mock_intake,
"work_item_attachments": mock_attachments,
"work_item_comments": mock_comments,
"work_item_links": mock_links,
"work_item_properties": mock_properties,
"work_item_types": mock_types,
"work_item_worklogs": mock_worklogs,
}
@pytest.fixture
def sdk_adapter_with_mock_apis(self, mock_sdk_apis):
"""Create a PlaneSDKAdapter with mocked SDK APIs."""
adapter = PlaneSDKAdapter(access_token="test-token")
# Replace the SDK API objects with mocks
adapter.intake = mock_sdk_apis["intake"]
adapter.work_item_attachments = mock_sdk_apis["work_item_attachments"]
adapter.work_item_comments = mock_sdk_apis["work_item_comments"]
adapter.work_item_links = mock_sdk_apis["work_item_links"]
adapter.work_item_properties = mock_sdk_apis["work_item_properties"]
adapter.work_item_types = mock_sdk_apis["work_item_types"]
adapter.work_item_worklogs = mock_sdk_apis["work_item_worklogs"]
return adapter
@pytest.mark.asyncio
async def test_intake_create_request_object_creation(self, sdk_adapter_with_mock_apis):
"""Test that intake create method properly creates request objects."""
adapter = sdk_adapter_with_mock_apis
# Call the method with kwargs - need to provide the correct nested structure
result = await adapter.create_intake_work_item(
workspace_slug="test-workspace",
project_id="project-123",
issue={"name": "Test Intake", "description_html": "<p>Test description</p>", "priority": "medium"},
intake="test-intake-id",
)
assert result["id"] == "intake-123"
# Verify the SDK method was called with the correct request object
adapter.intake.create_intake_work_item.assert_called_once()
call_args = adapter.intake.create_intake_work_item.call_args
# Check that the request object was created and passed correctly
assert call_args[1]["project_id"] == "project-123"
assert call_args[1]["slug"] == "test-workspace"
assert "intake_issue_create_request" in call_args[1]
# The request object should be an instance of IntakeIssueCreateRequest
request_obj = call_args[1]["intake_issue_create_request"]
assert hasattr(request_obj, "name")
assert hasattr(request_obj, "description_html")
assert hasattr(request_obj, "priority")
@pytest.mark.asyncio
async def test_intake_update_request_object_creation(self, sdk_adapter_with_mock_apis):
"""Test that intake update method properly creates patched request objects."""
adapter = sdk_adapter_with_mock_apis
# Call the method with kwargs
result = await adapter.update_intake_work_item(
workspace_slug="test-workspace", project_id="project-123", intake_id="intake-123", name="Updated Intake", priority="high"
)
assert result["id"] == "intake-123"
# Verify the SDK method was called with the correct patched request object
adapter.intake.update_intake_work_item.assert_called_once()
call_args = adapter.intake.update_intake_work_item.call_args
# Check that the patched request object was created and passed correctly
assert call_args[1]["pk"] == "intake-123"
assert call_args[1]["project_id"] == "project-123"
assert call_args[1]["slug"] == "test-workspace"
assert "patched_intake_issue_update_request" in call_args[1]
# The request object should be an instance of PatchedIntakeIssueUpdateRequest
request_obj = call_args[1]["patched_intake_issue_update_request"]
assert hasattr(request_obj, "name")
assert hasattr(request_obj, "priority")
@pytest.mark.asyncio
async def test_attachments_create_request_object_creation(self, sdk_adapter_with_mock_apis):
"""Test that attachments create method properly creates request objects."""
adapter = sdk_adapter_with_mock_apis
# Call the method with kwargs - need to provide required fields
result = await adapter.create_work_item_attachment(
workspace_slug="test-workspace", project_id="project-123", issue_id="issue-123", name="test-file.pdf", size=1024
)
assert result["id"] == "attachment-123"
# Verify the SDK method was called with the correct request object
adapter.work_item_attachments.create_work_item_attachment.assert_called_once()
call_args = adapter.work_item_attachments.create_work_item_attachment.call_args
# Check that the request object was created and passed correctly
assert call_args[1]["project_id"] == "project-123"
assert call_args[1]["slug"] == "test-workspace"
assert "issue_attachment_upload_request" in call_args[1]
# The request object should be an instance of IssueAttachmentUploadRequest
request_obj = call_args[1]["issue_attachment_upload_request"]
assert hasattr(request_obj, "name")
assert hasattr(request_obj, "size")
@pytest.mark.asyncio
async def test_comments_create_request_object_creation(self, sdk_adapter_with_mock_apis):
"""Test that comments create method properly creates request objects."""
adapter = sdk_adapter_with_mock_apis
# Call the method with kwargs
result = await adapter.create_work_item_comment(
workspace_slug="test-workspace", project_id="project-123", issue_id="issue-123", comment_html="<p>Test comment</p>"
)
assert result["id"] == "comment-123"
# Verify the SDK method was called with the correct request object
adapter.work_item_comments.create_work_item_comment.assert_called_once()
call_args = adapter.work_item_comments.create_work_item_comment.call_args
# Check that the request object was created and passed correctly
assert call_args[1]["project_id"] == "project-123"
assert call_args[1]["slug"] == "test-workspace"
assert "issue_comment_create_request" in call_args[1]
# The request object should be an instance of IssueCommentCreateRequest
request_obj = call_args[1]["issue_comment_create_request"]
assert hasattr(request_obj, "comment_html")
@pytest.mark.asyncio
async def test_comments_update_request_object_creation(self, sdk_adapter_with_mock_apis):
"""Test that comments update method properly creates patched request objects."""
adapter = sdk_adapter_with_mock_apis
# Call the method with kwargs
result = await adapter.update_work_item_comment(
workspace_slug="test-workspace", project_id="project-123", comment_id="comment-123", comment_html="<p>Updated comment</p>"
)
assert result["id"] == "comment-123"
# Verify the SDK method was called with the correct patched request object
adapter.work_item_comments.update_work_item_comment.assert_called_once()
call_args = adapter.work_item_comments.update_work_item_comment.call_args
# Check that the patched request object was created and passed correctly
assert call_args[1]["pk"] == "comment-123"
assert call_args[1]["project_id"] == "project-123"
assert call_args[1]["slug"] == "test-workspace"
assert "patched_issue_comment_create_request" in call_args[1]
# The request object should be an instance of PatchedIssueCommentCreateRequest
request_obj = call_args[1]["patched_issue_comment_create_request"]
assert hasattr(request_obj, "comment_html")
@pytest.mark.asyncio
async def test_worklogs_create_request_object_creation(self, sdk_adapter_with_mock_apis):
"""Test that worklogs create method properly creates request objects."""
adapter = sdk_adapter_with_mock_apis
# Call the method with kwargs
result = await adapter.create_issue_worklog(
workspace_slug="test-workspace", project_id="project-123", issue_id="issue-123", description="Worked on feature", duration=120
)
assert result["id"] == "worklog-123"
# Verify the SDK method was called with the correct request object
adapter.work_item_worklogs.create_issue_worklog.assert_called_once()
call_args = adapter.work_item_worklogs.create_issue_worklog.call_args
# Check that the request object was created and passed correctly
assert call_args[1]["project_id"] == "project-123"
assert call_args[1]["slug"] == "test-workspace"
assert "issue_work_log_api_request" in call_args[1]
# The request object should be an instance of IssueWorkLogAPIRequest
request_obj = call_args[1]["issue_work_log_api_request"]
assert hasattr(request_obj, "description")
assert hasattr(request_obj, "duration")
@pytest.mark.asyncio
async def test_all_request_object_types_covered(self, sdk_adapter_with_mock_apis):
"""Test that all the new request object types are properly handled."""
adapter = sdk_adapter_with_mock_apis
# Test all the different request object types
test_cases = [
# (method, kwargs, expected_request_param)
(
adapter.create_work_item_link,
{"workspace_slug": "test-workspace", "project_id": "project-123", "issue_id": "issue-123", "url": "https://example.com"},
"issue_link_create_request",
),
(
adapter.update_issue_link,
{"workspace_slug": "test-workspace", "project_id": "project-123", "link_id": "link-123", "url": "https://updated.com"},
"patched_issue_link_update_request",
),
(
adapter.create_issue_property,
{"workspace_slug": "test-workspace", "project_id": "project-123", "name": "Priority", "description": "Custom priority field"},
"issue_property_api_request",
),
(
adapter.update_issue_property,
{"workspace_slug": "test-workspace", "project_id": "project-123", "property_id": "property-123", "name": "Updated Priority"},
"patched_issue_property_api_request",
),
(
adapter.create_issue_type,
{"workspace_slug": "test-workspace", "project_id": "project-123", "name": "Bug", "description": "Bug report type"},
"issue_type_api_request",
),
(
adapter.update_issue_type,
{"workspace_slug": "test-workspace", "project_id": "project-123", "type_id": "type-123", "name": "Updated Bug"},
"patched_issue_type_api_request",
),
(
adapter.update_issue_worklog,
{
"workspace_slug": "test-workspace",
"project_id": "project-123",
"worklog_id": "worklog-123",
"description": "Updated work",
"duration": 180,
},
"patched_issue_work_log_api_request",
),
]
for method, kwargs, expected_request_param in test_cases:
# Call the method
result = await method(**kwargs)
# Verify the result is successful
assert result is not None
# Get the mock call
mock_call = method.call_args
assert mock_call is not None
# Verify the request object parameter exists
assert expected_request_param in mock_call[1], f"Expected {expected_request_param} in call args for {method.__name__}"
# Verify the request object has the expected attributes
request_obj = mock_call[1][expected_request_param]
assert request_obj is not None, f"Request object should not be None for {method.__name__}"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
apps/pi/pi/tests/verify_lazy_llm.py
+79
View File
@@ -0,0 +1,79 @@
import asyncio
import os
import sys
# Ensure we can import from the app
sys.path.append(os.getcwd())
# load_dotenv()
from pi.services.llm import llms
async def main():
print("Verifying LazyLLM implementation...")
# Access the global llm instance
lazy_instance = llms.llm
print(f"Global llm type: {type(lazy_instance)}")
print(f"Is instance of LazyLLM? {isinstance(lazy_instance, llms.LazyLLM)}")
# Check internal state BEFORE access
if hasattr(lazy_instance, "_proxy_target"):
print(f"Internal _proxy_target before access: {lazy_instance._proxy_target}")
if lazy_instance._proxy_target is None:
print("SUCCESS: LLM is not yet initialized.")
else:
print("FAILURE: LLM is already initialized!")
else:
print("FAILURE: LazyLLM structure not found.")
# Trigger initialization
print("\nTriggering initialization via attribute access...")
_ = lazy_instance.model_name
# Check internal state AFTER access
print(f"Internal _proxy_target after access: {lazy_instance._proxy_target}")
if lazy_instance._proxy_target is not None:
print("SUCCESS: LLM initialized on demand.")
else:
print("FAILURE: LLM did not initialize!")
print("\nVerification Complete for Global LLM.")
print("\nVerifying ChatKit.tool_llm implementation...")
from pi.services.chat.kit import ChatKit
# Instantiate ChatKit
chat_kit = ChatKit()
lazy_tool_llm = chat_kit.tool_llm
print(f"ChatKit.tool_llm type: {type(lazy_tool_llm)}")
print(f"Is instance of LazyLLM? {isinstance(lazy_tool_llm, llms.LazyLLM)}")
# Check internal state BEFORE access
if hasattr(lazy_tool_llm, "_proxy_target"):
print(f"Internal _proxy_target before access: {lazy_tool_llm._proxy_target}")
if lazy_tool_llm._proxy_target is None:
print("SUCCESS: ChatKit.tool_llm is not yet initialized.")
else:
print("FAILURE: ChatKit.tool_llm is already initialized!")
else:
print("FAILURE: LazyLLM structure not found in ChatKit.tool_llm.")
# Trigger initialization
print("\nTriggering ChatKit.tool_llm initialization via attribute access...")
_ = lazy_tool_llm.model_name
# Check internal state AFTER access
print(f"Internal _proxy_target after access: {lazy_tool_llm._proxy_target}")
if lazy_tool_llm._proxy_target is not None:
print("SUCCESS: ChatKit.tool_llm initialized on demand.")
else:
print("FAILURE: ChatKit.tool_llm did not initialize!")
print("\nVerification Complete.")
if __name__ == "__main__":
asyncio.run(main())
apps/pi/requirements.txt
+5 -5
View File
@@ -2,8 +2,8 @@
fastapi==0.111.1
requests==2.32.5
pydantic==2.8.2
plane-sdk==0.1.10
openai==1.77.0
plane-sdk==0.2.0
openai==2.8.1
html2text==2024.2.26
python-json-logger==3.3.0
sqlglot==27.2.0
@@ -16,8 +16,8 @@ boto3==1.40.16
deepgram-sdk==4.8.1
# langchain
langchain_openai==0.3.19
langchain_community==0.3.23
langchain_openai==1.1.0
langchain_community==0.4.1
# sentry
sentry-sdk==2.17.0
@@ -65,4 +65,4 @@ pytest==8.3.2
# markdown to html
markdown-it-py==4.0.0
mdit-py-plugins==0.5.0
mdit-py-plugins==0.5.0
apps/pi/tests/test_v1_v2_comparison.py
+2 -2
View File
@@ -1396,7 +1396,7 @@ class APITester:
v1_url=v1_create_url,
v2_url=v2_create_url,
status=TestStatus.PASS if error_match else TestStatus.FAIL,
message=f"Negative test: Both rejected invalid file during upload(V1: 400, V2: 400) - Error match: {'✓' if error_match else '✗'}",
message=f"Negative test: Both rejected invalid file during upload (V1: 400, V2: 400) - Error match: {'✓' if error_match else '✗'}", #noqa: E501
differences=None if error_match else {"error_messages": {"v1": v1_error_msg, "v2": v2_error_msg}},
v1_response=v1_create_data,
v2_response=v2_create_data,
@@ -1413,7 +1413,7 @@ class APITester:
v1_url=v1_create_url,
v2_url=v2_create_url,
status=TestStatus.FAIL,
message=f"Inconsistent behavior during upload: V1: {v1_create_resp.status_code}, V2: {v2_create_resp.status_code}",
message=f"Inconsistent behavior during upload: V1: {v1_create_resp.status_code}, V2: {v2_create_resp.status_code}", #noqa: E501
differences={"status_mismatch": {"v1": v1_create_resp.status_code, "v2": v2_create_resp.status_code}},
v1_response=v1_create_data,
v2_response=v2_create_data,
apps/web/ee/components/pi-chat/conversation/ai-message.tsx
+3 -3
View File
@@ -31,14 +31,14 @@ export const AiMessage = observer((props: TProps) => {
const { getWorkspaceBySlug } = useWorkspace();
// derived
const workspaceId = getWorkspaceBySlug(workspaceSlug as string)?.id;
const { query_id, answer, reasoning, isPiThinking, feedback } = dialogue || {};
const { query_id, answer, reasoning, current_tick, isPiThinking, feedback } = dialogue || {};
return (
<div className="flex gap-4" id={id}>
<div className="flex flex-col text-base break-words w-full">
{/* Message */}
<div className="flex flex-col gap-4">
{!isLoading && <ReasoningBlock reasoning={reasoning} isThinking={isPiThinking} />}
{!isLoading && <ReasoningBlock reasoning={reasoning} currentTick={current_tick} isThinking={isPiThinking} />}
<Markdown
remarkPlugins={[remarkGfm]}
className="pi-chat-root [&>*:first-child]:mt-0 animate-fade-in"
@@ -53,7 +53,7 @@ export const AiMessage = observer((props: TProps) => {
<table className="min-w-full border-collapse">{children}</table>
</div>
),
th: ({ children }) => <th className="px-2 py-3 border-custom-border-200">{children}</th>,
th: ({ children }) => <th className="px-2 py-3 border-custom-border-200"> {children}</th>,
td: ({ children }) => <td className="px-2 py-3 border-custom-border-200">{children}</td>,
}}
>
apps/web/ee/components/pi-chat/conversation/feedback.tsx
+1 -1
View File
@@ -62,7 +62,7 @@ export const Feedback = observer((props: TProps) => {
};
return (
<div className="flex gap-4 mt-6">
<div className="flex gap-4">
{/* Copy */}
<Tooltip tooltipContent="Copy to clipboard" position="bottom" className="mb-4">
<Copy size={16} onClick={handleCopyLink} className="my-auto cursor-pointer text-custom-text-300" />
apps/web/ee/components/pi-chat/conversation/reasoning.tsx
+13 -7
View File
@@ -1,16 +1,19 @@
import { useEffect, useState } from "react";
import { useState } from "react";
import Markdown from "react-markdown";
import { Brain, ChevronDownIcon } from "lucide-react";
import { cn } from "@plane/utils";
const stripEmojis = (str: string) => str.replace(/\p{Emoji}/gu, "");
type TProps = {
reasoning: string | undefined;
currentTick?: string;
reasoning?: string;
isThinking: boolean | undefined;
};
export const ReasoningBlock = (props: TProps) => {
const { reasoning, isThinking } = props;
const [isOpen, setIsOpen] = useState(true);
const { currentTick, reasoning, isThinking } = props;
const [isOpen, setIsOpen] = useState(false);
return (
<div className={cn("flex flex-col")}>
@@ -37,11 +40,14 @@ export const ReasoningBlock = (props: TProps) => {
onClick={() => setIsOpen(!isOpen)}
aria-expanded={isOpen}
className={cn(
"flex items-center gap-2 w-full px-3 transition-all duration-500 ease-in-out hover:border-transparent py-2"
"flex justify-start items-center gap-2 w-full px-3 transition-all duration-500 ease-in-out hover:border-transparent py-2 max-w-full overflow-hidden"
)}
>
<div className="w-2 h-4 rounded-[1px] pi-cursor animate-vertical-scale" />
<div className="flex gap-2 items-center shimmer">Thinking</div>
<div className={cn("shimmer truncate")}>{stripEmojis(currentTick || "Thinking")}</div>{" "}
<ChevronDownIcon
className={`w-4 h-4 transition-transform duration-500 ease-in-out flex-shrink-0 ${isOpen ? "transform rotate-180" : ""}`}
/>
</button>
)}
@@ -50,7 +56,7 @@ export const ReasoningBlock = (props: TProps) => {
"overflow-hidden",
"transition-all duration-500 ease-in-out",
isOpen ? "max-h-screen opacity-100" : "max-h-0 opacity-0 mt-0"
isOpen ? "max-h-fit opacity-100" : "max-h-0 opacity-0 mt-0"
)}
>
<div className="mx-4 overflow-hidden text-custom-text-300 relative">
apps/web/ee/components/pi-chat/converse/voice-input.tsx
+15 -4
View File
@@ -19,14 +19,25 @@ type TProps = {
isFullScreen: boolean;
focus: TFocus;
loader: TPiLoaders;
mode: string;
setLoader: Dispatch<SetStateAction<TPiLoaders>>;
createNewChat: (focus: TFocus, isProjectLevel: boolean, workspaceId: string) => Promise<string>;
createNewChat: (focus: TFocus, mode: string, isProjectLevel: boolean, workspaceId: string) => Promise<string>;
};
const piChatService = new PiChatService();
const AudioRecorder = (props: TProps) => {
const { workspaceId, chatId, editorRef, createNewChat, isProjectLevel, isFullScreen, focus, loader, setLoader } =
props;
const {
workspaceId,
chatId,
editorRef,
createNewChat,
isProjectLevel,
isFullScreen,
focus,
loader,
setLoader,
mode,
} = props;
const [isSubmitting, setIsSubmitting] = useState(false);
const [waveformData, setWaveformData] = useState<{ index: number; amplitude: number }[]>([]);
@@ -111,7 +122,7 @@ const AudioRecorder = (props: TProps) => {
try {
setLoader("transcribing");
let chatIdToUse = chatId;
if (!chatIdToUse) chatIdToUse = await createNewChat(focus, isProjectLevel, workspaceId);
if (!chatIdToUse) chatIdToUse = await createNewChat(focus, mode, isProjectLevel, workspaceId);
const response = await piChatService.transcribeAudio(workspaceId, formData, chatIdToUse);
editorRef.current?.appendText(" " + response);
} catch (err) {
apps/web/ee/components/pi-chat/input/dnd-wrapper.tsx
+4 -2
View File
@@ -24,7 +24,8 @@ type Props = {
chatId: string | undefined;
isProjectLevel: boolean;
focus: TFocus;
createNewChat: (focus: TFocus, isProjectLevel: boolean, workspaceId: string) => Promise<string>;
mode: string;
createNewChat: (focus: TFocus, mode: string, isProjectLevel: boolean, workspaceId: string) => Promise<string>;
setAttachments: Dispatch<SetStateAction<TPiAttachment[]>>;
children: (isUploading: boolean, open: () => void) => React.ReactNode;
};
@@ -39,6 +40,7 @@ export const DndWrapper: FC<Props> = observer((props) => {
isProjectLevel,
createNewChat,
focus,
mode,
children,
} = props;
@@ -62,7 +64,7 @@ export const DndWrapper: FC<Props> = observer((props) => {
if (rejectedFiles.length === 0) {
setIsUploading(true);
let chatIdToUse = chatId;
if (!chatIdToUse) chatIdToUse = await createNewChat(focus, isProjectLevel, workspaceId);
if (!chatIdToUse) chatIdToUse = await createNewChat(focus, mode, isProjectLevel, workspaceId);
for (const file of acceptedFiles) {
const currentFile: File = file;
if (!currentFile) return;
apps/web/ee/components/pi-chat/input/focus-filter.tsx
+4 -4
View File
@@ -81,7 +81,7 @@ export const FocusFilter = observer((props: TProps) => {
className="ml-4 max-w-[200px] font-medium text-custom-text-300"
disabled={focus.isInWorkspaceContext}
>
<div className="flex rounded-full font-medium gap-2 w-full overflow-hidden">
<div className="flex font-medium gap-2 w-full overflow-hidden">
{!isEmpty(focus) && !isEmpty(selectedFocus) && focus.isInWorkspaceContext ? (
<div className=" flex items-center gap-2 text-sm my-auto capitalize truncate">
{focus.entityType === "workspace_id" ? (
@@ -96,8 +96,8 @@ export const FocusFilter = observer((props: TProps) => {
<span className="truncate">{selectedFocus?.name}</span>
</div>
) : (
<div className="text-sm font-medium text-custom-text-300 my-auto flex items-center gap-2">
<AtSign className="size-4" /> <span>Add context </span>
<div className="text-sm font-medium my-auto flex items-center gap-2">
<AtSign className="size-4 text-custom-text-300" /> <span>Add context </span>
</div>
)}
</div>
@@ -114,7 +114,7 @@ export const FocusFilter = observer((props: TProps) => {
maxHeight="lg"
className="flex flex-col-reverse"
buttonClassName={cn(
"rounded-md h-full px-2 border-[0.5px] border-custom-border-200 max-h-[30px] overflow-hidden max-w-[200px] hover:bg-custom-background-100"
"rounded-lg h-full px-2 border-[0.5px] border-custom-border-200 max-h-[30px] overflow-hidden max-w-[200px] hover:bg-custom-background-100"
)}
optionsClassName="max-h-[70vh] overflow-y-auto"
>
apps/web/ee/components/pi-chat/input/mode.tsx
+61
View File
@@ -0,0 +1,61 @@
import { AiSearchIcon, AiWriteIcon } from "@plane/propel/icons";
import { CustomSelect, Tooltip } from "@plane/ui";
import { cn } from "@plane/utils";
const MODE_OPTIONS = [
{
label: "Ask",
value: "ask",
icon: <AiSearchIcon className="size-4" />,
},
{
label: "Build",
value: "build",
icon: <AiWriteIcon className="size-4" />,
},
];
export const AiMode = (props: { aiMode: string; setAiMode: (mode: string) => void }) => {
const { aiMode, setAiMode } = props;
const selectedMode = MODE_OPTIONS.find((option) => option.value === aiMode);
return (
<CustomSelect
value={aiMode}
label={
<Tooltip tooltipContent="Select the mode of the AI to use for the conversation." position="top">
<div className="flex font-medium gap-2 w-full overflow-hidden">
<span className="flex items-center gap-2 text-custom-text-300">{selectedMode?.icon}</span>
<span className="text-sm truncate">{selectedMode?.label}</span>
</div>
</Tooltip>
}
noChevron
onChange={(val: string) => {
setAiMode(val);
}}
maxHeight="lg"
className="flex flex-col-reverse"
buttonClassName={cn(
"rounded-lg h-full px-2 border-[0.5px] border-custom-border-200 h-[27px] overflow-hidden hover:bg-custom-background-100"
)}
optionsClassName="max-h-[70vh] overflow-y-auto"
>
<div className="flex flex-col divide-y divide-custom-border-100 space-y-2 max-h-full">
<div>
{MODE_OPTIONS.map((option) => (
<CustomSelect.Option
key={option.value}
value={option.value}
className="text-sm text-custom-text-200 font-medium flex w-full"
>
<div className="flex items-center gap-2">
<span className="flex items-center gap-2 text-custom-text-300">{option.icon}</span>
<span className="text-sm truncate">{option.label}</span>
</div>
</CustomSelect.Option>
))}
</div>
</div>
</CustomSelect>
);
};
apps/web/ee/components/pi-chat/input/root.tsx
+33 -19
View File
@@ -24,6 +24,7 @@ import { formatSearchQuery } from "../helper";
import { InputPreviewUploads } from "../uploads/input-preview-uploads";
import { DndWrapper } from "./dnd-wrapper";
import { FocusFilter } from "./focus-filter";
import { AiMode } from "./mode";
import { AttachmentActionButton } from "./quick-action-button";
type TEditCommands = {
@@ -61,6 +62,7 @@ export const InputBox = observer((props: TProps) => {
getChatFocus,
fetchModels,
abortStream,
getChatMode,
attachmentStore: { getAttachmentsUploadStatusByChatId },
} = usePiChat();
const { getWorkspaceBySlug } = useWorkspace();
@@ -76,6 +78,7 @@ export const InputBox = observer((props: TProps) => {
const projectDetails = getProjectByIdentifier(projectIdentifier);
const projectIdToUse = projectDetails?.id || projectId || "";
const chatFocus = getChatFocus(activeChatId);
const chatMode = getChatMode(activeChatId || "");
const attachmentsUploadStatus = getAttachmentsUploadStatusByChatId(activeChatId || "");
// state
const [focus, setFocus] = useState<TFocus>(
@@ -87,6 +90,7 @@ export const InputBox = observer((props: TProps) => {
);
const [loader, setLoader] = useState<TPiLoaders>("");
const [attachments, setAttachments] = useState<TPiAttachment[]>([]);
const [aiMode, setAiMode] = useState<string>(chatMode);
const [isEditorReady, setIsEditorReady] = useState(false);
//ref
const editorCommands = useRef<TEditCommands | null>(null);
@@ -126,7 +130,7 @@ export const InputBox = observer((props: TProps) => {
let chatIdToUse = activeChatId;
setLoader("submitting");
if (!chatIdToUse) {
chatIdToUse = await createNewChat(focus, isProjectLevel, workspaceId);
chatIdToUse = await createNewChat(focus, aiMode, isProjectLevel, workspaceId);
}
// Don't redirect if we are in the floating chat window
if (shouldRedirect && !routeChatId)
@@ -142,7 +146,8 @@ export const InputBox = observer((props: TProps) => {
workspaceSlug?.toString(),
workspaceId,
pathname,
attachmentIds
attachmentIds,
aiMode
);
editorCommands.current?.clear();
addContext();
@@ -169,6 +174,9 @@ export const InputBox = observer((props: TProps) => {
};
setFocus(presentFocus);
}
if (chatMode) {
setAiMode(chatMode);
}
}, [isChatLoading, chatFocus]);
// Adding context for the sidecar
@@ -201,6 +209,7 @@ export const InputBox = observer((props: TProps) => {
chatId={activeChatId}
setAttachments={setAttachments}
isProjectLevel={isProjectLevel}
mode={aiMode}
createNewChat={createNewChat}
focus={focus}
>
@@ -223,22 +232,9 @@ export const InputBox = observer((props: TProps) => {
/>
</div>
)}
{/* editor view */}
<PiChatEditorWithRef
setEditorCommand={(command) => {
setEditorCommands({ ...command });
}}
handleSubmit={handleSubmit}
searchCallback={getMentionSuggestions}
className={cn("flex-1 max-h-[250px] min-h-[70px]", {
"absolute w-0": SPEECH_LOADERS.includes(loader),
})}
onEditorReady={() => setIsEditorReady(true)}
ref={editorRef}
/>
<div className="flex w-full gap-3 justify-between">
{/* Focus */}
{!SPEECH_LOADERS.includes(loader) && (
{/* Focus */}
{!SPEECH_LOADERS.includes(loader) && (
<div className="mb-2 w-fit">
<FocusFilter
workspaceId={workspaceId}
projectId={projectIdToUse}
@@ -246,7 +242,24 @@ export const InputBox = observer((props: TProps) => {
setFocus={setFocus}
isLoading={isChatLoading && !!activeChatId}
/>
)}
</div>
)}
{/* editor view */}
<PiChatEditorWithRef
setEditorCommand={(command) => {
setEditorCommands({ ...command });
}}
handleSubmit={handleSubmit}
searchCallback={getMentionSuggestions}
className={cn("flex-1 max-h-[250px] min-h-[50px]", {
"absolute w-0": SPEECH_LOADERS.includes(loader),
})}
onEditorReady={() => setIsEditorReady(true)}
ref={editorRef}
/>
<div className="flex w-full gap-3 justify-between">
{/* Focus */}
{!SPEECH_LOADERS.includes(loader) && <AiMode aiMode={aiMode} setAiMode={setAiMode} />}
<div className="flex items-center w-full justify-end gap-2">
<div className="flex w-full justify-end">
{/* speech recorder */}
@@ -265,6 +278,7 @@ export const InputBox = observer((props: TProps) => {
setLoader={setLoader}
isFullScreen={isFullScreen}
focus={focus}
mode={aiMode}
/>
</WithFeatureFlagHOC>
{!SPEECH_LOADERS.includes(loader) && (
apps/web/ee/components/pi-chat/system-prompts.tsx
+3 -3
View File
@@ -50,7 +50,7 @@ const SystemPrompts = (props: TSystemPrompt) => {
entityType: projectId ? "project_id" : "workspace_id",
entityIdentifier: projectId?.toString() || workspaceId?.toString() || "",
};
const newChatId = await createNewChat(focus, isProjectLevel, workspaceId);
const newChatId = await createNewChat(focus, "ask", isProjectLevel, workspaceId);
setIsInitializing("");
// Don't redirect if we are in the floating chat window
if (shouldRedirect) router.push(`/${workspaceSlug}/${isProjectLevel ? "projects/" : ""}pi-chat/${newChatId}`);
@@ -62,7 +62,8 @@ const SystemPrompts = (props: TSystemPrompt) => {
workspaceSlug?.toString(),
workspaceId?.toString(),
pathname,
[]
[],
"ask"
);
};
const promptIcon = getIcon(prompt.type);
@@ -92,7 +93,6 @@ const SystemPrompts = (props: TSystemPrompt) => {
</span>
)}
</div>
<span className="text-left text-sm break-words line-clamp-2">{prompt.text}</span>
</button>
);
apps/web/ee/store/pi-chat/pi-chat.ts
+36 -7
View File
@@ -51,6 +51,7 @@ export interface IPiChatStore {
geFavoriteChats: () => TUserThreads[];
getChatById: (chatId: string) => TChatHistory;
getChatFocus: (chatId: string | undefined) => TFocus | undefined;
getChatMode: (chatId: string) => string;
// actions
initPiChat: (chatId?: string) => void;
fetchChatById: (chatId: string, workspaceId: string | undefined) => void;
@@ -62,7 +63,8 @@ export interface IPiChatStore {
workspaceSlug: string,
workspaceId: string | undefined,
callbackUrl: string,
attachmentIds: string[]
attachmentIds: string[],
aiMode: string
) => Promise<void>;
getInstance: (workspaceId: string) => Promise<TInstanceResponse>;
fetchUserThreads: (workspaceId: string | undefined, isProjectChat: boolean) => void;
@@ -78,7 +80,12 @@ export interface IPiChatStore {
fetchModels: (workspaceId?: string) => Promise<void>;
abortStream: (chatId: string) => void;
setActiveModel: (model: TAiModels) => void;
createNewChat: (focus: TFocus, isProjectChat: boolean, workspaceId: string | undefined) => Promise<string>;
createNewChat: (
focus: TFocus,
mode: string,
isProjectChat: boolean,
workspaceId: string | undefined
) => Promise<string>;
renameChat: (chatId: string, title: string, workspaceId: string | undefined) => Promise<void>;
deleteChat: (chatId: string, workspaceSlug: string) => Promise<void>;
favoriteChat: (chatId: string, workspaceId: string | undefined) => Promise<void>;
@@ -243,6 +250,14 @@ export class PiChatStore implements IPiChatStore {
}
});
getChatMode = computedFn((chatId: string) => {
const chat = this.chatMap[chatId];
if (chat) {
return chat.mode || "ask";
}
return "ask";
});
getGroupedArtifactsByDialogue = computedFn((chatId: string, messageId: string) => {
const dialogue = this.chatMap[chatId]?.dialogueMap[messageId];
const response: {
@@ -294,7 +309,12 @@ export class PiChatStore implements IPiChatStore {
this.isPiArtifactsDrawerOpen = undefined;
};
createNewChat = async (focus: TFocus, isProjectChat: boolean = false, workspaceId: string | undefined) => {
createNewChat = async (
focus: TFocus,
mode: string,
isProjectChat: boolean = false,
workspaceId: string | undefined
) => {
this.isNewChat = true;
let payload: TInitPayload = {
workspace_in_context: focus.isInWorkspaceContext,
@@ -323,6 +343,7 @@ export class PiChatStore implements IPiChatStore {
focus_workspace_id: focus.entityType === "workspace_id" ? focus.entityIdentifier : "",
focus_project_id: focus.entityType === "project_id" ? focus.entityIdentifier : "",
workspace_id: workspaceId,
mode,
};
if (isProjectChat) {
@@ -343,7 +364,8 @@ export class PiChatStore implements IPiChatStore {
workspaceId: string | undefined,
callbackUrl: string | undefined,
attachmentIds: string[],
isNewChat: boolean
isNewChat: boolean,
mode: string
) => {
let payload: TQuery = {
chat_id: chatId,
@@ -364,6 +386,7 @@ export class PiChatStore implements IPiChatStore {
workspace_slug: workspaceSlug,
workspace_id: workspaceId,
attachment_ids: attachmentIds,
mode,
};
if (focus.isInWorkspaceContext) {
payload = { ...payload, [focus.entityType]: focus.entityIdentifier };
@@ -410,7 +433,9 @@ export class PiChatStore implements IPiChatStore {
eventSource.addEventListener("reasoning", (event: MessageEvent) => {
try {
const data = JSON.parse(event.data);
newDialogue.reasoning += data.reasoning;
if (!data.header) return;
newDialogue.current_tick = data.header;
newDialogue.reasoning = newDialogue.reasoning + `${data.header}${data.content}`;
this.updateDialogue(chatId, token, newDialogue);
} catch (e) {
console.error("Reasoning parse error", e);
@@ -478,7 +503,8 @@ export class PiChatStore implements IPiChatStore {
workspaceSlug: string,
workspaceId: string | undefined,
callbackUrl: string,
attachmentIds: string[]
attachmentIds: string[],
aiMode: string
) => {
if (!chatId) {
throw new Error("Chat not initialized");
@@ -489,6 +515,7 @@ export class PiChatStore implements IPiChatStore {
query,
llm: this.activeModel?.id,
answer: "",
current_tick: "",
reasoning: "",
isPiThinking: true,
actions: [],
@@ -522,7 +549,8 @@ export class PiChatStore implements IPiChatStore {
workspaceId,
callbackUrl,
attachmentIds,
isNewChat
isNewChat,
aiMode
);
this.piChatService
.retrieveToken(payload)
@@ -578,6 +606,7 @@ export class PiChatStore implements IPiChatStore {
query: this.chatMap[chatId]?.dialogueMap?.[token]?.query,
llm: this.activeModel?.id,
answer: "",
current_tick: "",
reasoning: "",
isPiThinking: true,
query_id: token,
apps/web/ee/types/pi-chat.d.ts
Vendored
+3
View File
@@ -38,6 +38,7 @@ export type TQuery = {
pi_sidebar_open?: boolean;
sidebar_open_url?: string;
attachment_ids?: string[];
mode?: string;
};
export type TInitPayload = Pick<
TQuery,
@@ -138,6 +139,7 @@ export type TDialogue = {
answer?: string;
llm?: string;
feedback?: EFeedback;
current_tick?: string;
reasoning?: string;
isPiThinking: boolean;
execution_status?: EExecutionStatus;
@@ -163,6 +165,7 @@ export type TChatHistory = {
focus_workspace_id: string;
focus_project_id: string;
workspace_id?: string;
mode?: string;
};
export type TAction = {
packages/propel/src/icons/ai-search.tsx
+22
View File
@@ -0,0 +1,22 @@
import * as React from "react";
import type { ISvgIcons } from "./type";
export const AiSearchIcon: React.FC<ISvgIcons> = ({ className = "", ...rest }) => (
<svg
viewBox="0 0 16 16"
className={`${className}`}
stroke="transparent"
xmlns="http://www.w3.org/2000/svg"
fill="none"
{...rest}
>
<path
d="M14 14L11.1 11.1M7.33333 4C9.17428 4 10.6667 5.49238 10.6667 7.33333M12.6667 7.33333C12.6667 10.2789 10.2789 12.6667 7.33333 12.6667C4.38781 12.6667 2 10.2789 2 7.33333C2 4.38781 4.38781 2 7.33333 2C10.2789 2 12.6667 4.38781 12.6667 7.33333Z"
stroke="currentColor"
strokeWidth="1.25"
strokeLinecap="round"
strokeLinejoin="round"
/>
</svg>
);
packages/propel/src/icons/ai-write.tsx
+23
View File
@@ -0,0 +1,23 @@
import * as React from "react";
import type { ISvgIcons } from "./type";
export const AiWriteIcon: React.FC<ISvgIcons> = ({ className = "", ...rest }) => (
<svg
viewBox="0 0 16 16"
className={`${className}`}
stroke="transparent"
xmlns="http://www.w3.org/2000/svg"
fill="none"
{...rest}
>
<path
d="M10.8912 1.55807C11.8717 0.577612 13.4615 0.577611 14.442 1.55807C15.4224 2.53853 15.4224 4.12841 14.442 5.10885L5.63729 13.9126C5.45891 14.0909 5.31066 14.2442 5.13241 14.3618C4.98526 14.4588 4.82547 14.535 4.6578 14.5893C4.45463 14.6551 4.24256 14.6757 3.99179 14.7036L1.73593 14.9546C1.54722 14.9755 1.35846 14.9091 1.22421 14.7749C1.09024 14.6406 1.02455 14.4526 1.0455 14.2641L1.29647 12.0083C1.32434 11.7575 1.34496 11.5454 1.41073 11.3422C1.46503 11.1746 1.54127 11.0148 1.63827 10.8676C1.75583 10.6894 1.90907 10.5402 2.08749 10.3618L10.8912 1.55807ZM13.5582 2.44186C13.0659 1.94956 12.2673 1.94956 11.775 2.44186L2.97128 11.2465C2.75795 11.4599 2.71373 11.5074 2.68221 11.5551C2.6471 11.6084 2.61886 11.6663 2.59921 11.727C2.58157 11.7815 2.57204 11.8456 2.53866 12.146L2.37362 13.6255L3.85409 13.4614C4.15439 13.428 4.21856 13.4185 4.27304 13.4008C4.33368 13.3812 4.39169 13.3529 4.44491 13.3178C4.49261 13.2863 4.54021 13.2421 4.7535 13.0288L13.5582 4.22506C14.0505 3.73278 14.0504 2.93416 13.5582 2.44186Z"
fill="currentColor"
/>
<path
d="M2.55194 1.16482C2.69105 0.803119 2.76061 0.62227 2.86102 0.569808C2.94805 0.524338 3.05183 0.524338 3.13886 0.569808C3.23927 0.62227 3.30883 0.803119 3.44795 1.16482L3.578 1.50296C3.672 1.74737 3.71901 1.86957 3.7921 1.97236C3.85688 2.06347 3.93648 2.14307 4.02758 2.20785C4.13037 2.28094 4.25258 2.32794 4.49699 2.42194L4.83513 2.552C5.19682 2.69111 5.37767 2.76067 5.43014 2.86108C5.47561 2.94811 5.47561 3.05189 5.43014 3.13892C5.37767 3.23934 5.19683 3.30889 4.83513 3.44801L4.49699 3.57806C4.25258 3.67207 4.13037 3.71907 4.02758 3.79216C3.93648 3.85694 3.85688 3.93654 3.7921 4.02764C3.71901 4.13043 3.672 4.25264 3.578 4.49705L3.44795 4.83519C3.30883 5.19689 3.23928 5.37773 3.13886 5.4302C3.05183 5.47567 2.94805 5.47567 2.86102 5.4302C2.76061 5.37773 2.69105 5.19689 2.55194 4.83519L2.42188 4.49705C2.32788 4.25264 2.28088 4.13043 2.20779 4.02764C2.14301 3.93654 2.06341 3.85694 1.9723 3.79216C1.86951 3.71907 1.74731 3.67207 1.5029 3.57806L1.16476 3.44801C0.803058 3.30889 0.622209 3.23934 0.569747 3.13892C0.524277 3.05189 0.524277 2.94811 0.569747 2.86108C0.622209 2.76067 0.803058 2.69111 1.16476 2.552L1.50289 2.42194C1.74731 2.32794 1.86951 2.28094 1.9723 2.20785C2.06341 2.14307 2.143 2.06347 2.20778 1.97236C2.28088 1.86957 2.32788 1.74737 2.42188 1.50296L2.55194 1.16482Z"
fill="currentColor"
/>
</svg>
);
packages/propel/src/icons/index.ts
+2
View File
@@ -5,6 +5,8 @@ export { ICON_REGISTRY } from "./registry";
export * from "./actions";
export * from "./arrows";
export * from "./ai-icon";
export * from "./ai-search";
export * from "./ai-write";
export * from "./approver-icon";
export * from "./at-risk-icon";
export * from "./attachments";