diff --git a/experiencemaker/module/context_generator/step_context_generator.py b/experiencemaker/module/context_generator/step_context_generator.py index 9223493c..3275f989 100644 --- a/experiencemaker/module/context_generator/step_context_generator.py +++ b/experiencemaker/module/context_generator/step_context_generator.py @@ -1,62 +1,41 @@ import json import re +from pathlib import Path from typing import List from loguru import logger -from pydantic import Field, model_validator +from pydantic import Field from experiencemaker.enumeration.role import Role -from experiencemaker.module.context_generator.base_context_generator import BaseContextGenerator +from experiencemaker.module.prompt.prompt_mixin import PromptMixin +from experiencemaker.module.context_generator.base_context_generator import BaseContextGenerator, \ +CONTEXT_GENERATOR_REGISTRY from experiencemaker.schema.trajectory import Trajectory, ContextMessage, Message from experiencemaker.schema.vector_store_node import VectorStoreNode -from experiencemaker.storage.es_vector_store import EsVectorStore -from experiencemaker.storage.file_vector_store import FileVectorStore -class StepContextGenerator(BaseContextGenerator): +@CONTEXT_GENERATOR_REGISTRY.register("step") +class StepContextGenerator(BaseContextGenerator, PromptMixin): """ Step-level context generator that retrieves and utilizes step-level experiences from the experience store to provide relevant context for agent execution """ - # Vector Store Configuration - vector_store_type: str = Field(default="file_vector_store") - vector_store_hosts: str | List[str] = Field(default="http://localhost:9200") - vector_store_index_name: str = Field(default="step_experience_store") - store_dir: str = Field(default="./step_experiences/") + # Retrieval Configuration - can be configured via startup parameters + vector_retrieve_top_k: int = Field(default=15, description="Number of candidates to retrieve from vector store") + final_top_k: int = Field(default=5, description="Final number of experiences to return") + min_score_threshold: float = Field(default=0.3, description="Minimum score threshold for filtering") - # Retrieval Configuration - vector_retrieve_top_k: int = Field(default=15) - final_top_k: int = Field(default=5) - min_score_threshold: float = Field(default=0.3) + # Feature Switches - can be configured via startup parameters + enable_llm_rerank: bool = Field(default=True, description="Enable LLM-based reranking") + enable_context_rewrite: bool = Field(default=True, description="Enable context rewriting") + enable_score_filter: bool = Field(default=False, description="Enable score-based filtering") - # Feature Switches - enable_llm_rerank: bool = Field(default=True) - enable_context_rewrite: bool = Field(default=True) - enable_score_filter: bool = Field(default=True) - - @model_validator(mode="after") - def init_vector_store(self): - """Initialize vector store based on configuration""" - if self.vector_store_type == "file_vector_store": - self.vector_store = FileVectorStore( - embedding_model=self.embedding_model, - index_name=self.vector_store_index_name, - store_dir=self.store_dir - ) - elif self.vector_store_type == "es_vector_store": - self.vector_store = EsVectorStore( - embedding_model=self.embedding_model, - index_name=self.vector_store_index_name, - hosts=self.vector_store_hosts - ) - else: - raise ValueError(f"Unknown vector store type: {self.vector_store_type}") - - return self + # Prompt configuration + prompt_file_path: Path = Field(default=Path(__file__).parent / "step_context_generator_prompt.yaml") def _build_retrieve_query(self, trajectory: Trajectory, **kwargs) -> str: - """Build retrieval query from trajectory""" + """Build retrieval query from trajectory (implements base class method)""" # Use the original query as base base_query = trajectory.query @@ -67,12 +46,90 @@ class StepContextGenerator(BaseContextGenerator): return base_query - def vector_retrieve(self, query: str, top_k: int = 10) -> List[VectorStoreNode]: - """Vector similarity retrieval from experience store""" + def _retrieve_by_query(self, trajectory: Trajectory, query: str, workspace_id: str, retrieve_top_k: int, + **kwargs) -> List[VectorStoreNode]: + """Retrieve experiences by query (implements base class method)""" if not query: logger.warning("Empty query provided for vector retrieval") return [] + try: + # Use hybrid retrieval strategy + return self._hybrid_retrieve(query, trajectory, retrieve_top_k) + + except Exception as e: + logger.error(f"Error in vector retrieval: {e}") + return [] + + def _generate_context_message(self, + trajectory: Trajectory, + nodes: List[VectorStoreNode], + **kwargs) -> ContextMessage: + """Generate context message from retrieved experiences (implements base class method)""" + if not nodes: + return ContextMessage(content="") + + try: + # Format retrieved experiences + formatted_experiences = self._format_experiences_for_context(nodes) + + prompt = self.prompt_format( + prompt_name="context_generation_prompt", + query=trajectory.query, + current_step=kwargs.get("current_step", ""), + retrieved_experiences=formatted_experiences, + num_experiences=len(nodes) + ) + + response = self.llm.chat([Message(role=Role.USER, content=prompt)]) + + # Extract generated context from response + context_content = self._parse_json_response(response.content, "context") + + if not context_content: + # Fallback to simple formatting + context_content = self._create_simple_context(nodes) + + # Optionally rewrite context to make it more relevant + if self.enable_context_rewrite: + context_content = self._rewrite_context(trajectory.query, context_content, trajectory) + + return ContextMessage(content=context_content) + + except Exception as e: + logger.error(f"Error generating context message: {e}") + return ContextMessage(content=self._create_simple_context(nodes)) + + def _hybrid_retrieve(self, query: str, trajectory: Trajectory, top_k: int) -> List[VectorStoreNode]: + """Hybrid retrieval strategy combining multiple approaches""" + logger.info(f"Starting hybrid retrieval for query: '{query}'") + + # Step 1: Vector retrieval to get candidates + candidates = self._vector_retrieve(query, self.vector_retrieve_top_k) + + if not candidates: + logger.warning("No candidates found in vector retrieval") + return [] + + # Step 2: LLM reranking (optional) + if self.enable_llm_rerank: + candidates = self._llm_rerank(query, candidates) + + # Step 3: Score-based filtering (optional) + if self.enable_score_filter: + candidates = self._score_based_filter(candidates, self.min_score_threshold) + + # Step 4: Return top-k results + final_results = candidates[:top_k] + logger.info(f"Hybrid retrieval completed: {len(final_results)} experiences selected") + + return final_results + + def _vector_retrieve(self, query: str, top_k: int) -> List[VectorStoreNode]: + """Vector similarity retrieval from experience store""" + if not query: + return [] + try: retrieved_nodes = self.vector_store.retrieve_by_query( query=query, @@ -85,16 +142,17 @@ class StepContextGenerator(BaseContextGenerator): logger.error(f"Error in vector retrieval: {e}") return [] - def llm_rerank(self, query: str, candidates: List[VectorStoreNode]) -> List[VectorStoreNode]: + def _llm_rerank(self, query: str, candidates: List[VectorStoreNode]) -> List[VectorStoreNode]: """LLM-based reranking of candidate experiences""" - if not self.enable_llm_rerank or not candidates: + if not candidates: return candidates try: # Format candidates for LLM evaluation candidates_text = self._format_candidates_for_rerank(candidates) - prompt = self.prompt_handler.experience_rerank_prompt.format( + prompt = self.prompt_format( + prompt_name="experience_rerank_prompt", query=query, candidates=candidates_text, num_candidates=len(candidates) @@ -119,16 +177,17 @@ class StepContextGenerator(BaseContextGenerator): logger.error(f"Error in LLM reranking: {e}") return candidates - def llm_rewrite_context(self, query: str, context_content: str, trajectory: Trajectory) -> str: + def _rewrite_context(self, query: str, context_content: str, trajectory: Trajectory) -> str: """LLM-based context rewriting to make experiences more relevant and actionable for current task""" - if not self.enable_query_rewrite or not context_content: + if not context_content: return context_content try: # Extract current trajectory context current_context = self._extract_trajectory_context(trajectory) - prompt = self.prompt_handler.context_rewrite_prompt.format( + prompt = self.prompt_format( + prompt_name="context_rewrite_prompt", current_query=query, current_context=current_context, original_context=context_content @@ -136,7 +195,7 @@ class StepContextGenerator(BaseContextGenerator): response = self.llm.chat([Message(role=Role.USER, content=prompt)]) - # Extract rewritten context from JSON + # Extract rewritten context rewritten_context = self._parse_json_response(response.content, "rewritten_context") if rewritten_context and rewritten_context.strip(): @@ -149,12 +208,8 @@ class StepContextGenerator(BaseContextGenerator): logger.error(f"Error in context rewriting: {e}") return context_content - def score_based_filter(self, experiences: List[VectorStoreNode], - min_score: float) -> List[VectorStoreNode]: + def _score_based_filter(self, experiences: List[VectorStoreNode], min_score: float) -> List[VectorStoreNode]: """Filter experiences based on quality scores""" - if not self.enable_score_filter: - return experiences - filtered_experiences = [] for exp in experiences: @@ -173,103 +228,6 @@ class StepContextGenerator(BaseContextGenerator): logger.info(f"Score filtering: {len(filtered_experiences)}/{len(experiences)} experiences retained") return filtered_experiences - def hybrid_retrieve(self, query: str, trajectory: Trajectory, top_k: int = 5) -> List[VectorStoreNode]: - """Hybrid retrieval strategy combining multiple approaches""" - logger.info(f"Starting hybrid retrieval for query: '{query}'") - - # Step 1: Vector retrieval to get candidates - candidates = self.vector_retrieve(query, self.vector_retrieve_top_k) - - if not candidates: - logger.warning("No candidates found in vector retrieval") - return [] - - # Step 2: LLM reranking (optional) - reranked = self.llm_rerank(query, candidates) - - # Step 3: Score-based filtering (optional) - filtered = self.score_based_filter(reranked, self.min_score_threshold) - - # Step 4: Return top-k results - final_results = filtered[:top_k] - logger.info(f"Hybrid retrieval completed: {len(final_results)} experiences selected") - - return final_results - - def retrieve_by_query(self, trajectory: Trajectory, query: str, **kwargs) -> List[VectorStoreNode]: - """Retrieve experiences by query (implements base class method)""" - return self.hybrid_retrieve(query, trajectory, self.final_top_k) - - def generate_context_message(self, - trajectory: Trajectory, - nodes: List[VectorStoreNode], - **kwargs) -> ContextMessage: - """Generate context message from retrieved experiences""" - if not nodes: - return ContextMessage(content="") - - try: - # Format retrieved experiences - formatted_experiences = self._format_experiences_for_context(nodes) - - prompt = self.prompt_handler.context_generation_prompt.format( - query=trajectory.query, - current_step=kwargs.get("current_step", ""), - retrieved_experiences=formatted_experiences, - num_experiences=len(nodes) - ) - - response = self.llm.chat([Message(role=Role.USER, content=prompt)]) - - # Extract generated context from JSON - context_content = self._parse_json_response(response.content, "context") - - if not context_content: - # Fallback to simple formatting - context_content = self._create_context(nodes) - - return ContextMessage(content=context_content) - - except Exception as e: - logger.error(f"Error generating context message: {e}") - return ContextMessage(content=self._create_context(nodes)) - - def build_context_messages(self, task: str, experiences: List[VectorStoreNode], trajectory: Trajectory) -> List[ - Message]: - """Build context messages from experiences for agent consumption""" - if not experiences: - return [] - - messages = [] - - # Create initial context content with experiences - system_content = "You have access to the following relevant experiences from previous executions:\n\n" - - for i, exp in enumerate(experiences, 1): - condition = exp.content - experience_content = exp.metadata.get("experience", "") - tags = exp.metadata.get("tags", []) - - system_content += f"**Experience {i}:**\n" - system_content += f"When to use: {condition}\n" - system_content += f"Experience: {experience_content}\n" - system_content += f"Tags: {', '.join(tags)}\n\n" - - system_content += "Consider these experiences when planning and executing your approach." - - # Rewrite the complete context to make it more relevant to current task - if self.enable_context_rewrite: - system_content = self.llm_rewrite_context(task, system_content, trajectory) - - messages.append(Message(role=Role.SYSTEM, content=system_content)) - - return messages - - def get_best_experiences(self, task: str, trajectory: Trajectory, max_count: int = 3) -> List[Message]: - """Get the best relevant experiences for a task as formatted messages""" - experiences = self.hybrid_retrieve(task, trajectory, max_count) - return self.build_context_messages(task, experiences, trajectory) - def _extract_trajectory_context(self, trajectory: Trajectory) -> str: """Extract relevant context from trajectory for query enhancement""" context_parts = [] @@ -286,7 +244,7 @@ class StepContextGenerator(BaseContextGenerator): context_parts.append("Recent steps:\n" + "\n".join(step_summaries)) # Add metadata if available - if trajectory.metadata: + if hasattr(trajectory, 'metadata') and trajectory.metadata: relevant_metadata = {k: v for k, v in trajectory.metadata.items() if k in ["domain", "task_type", "difficulty"]} if relevant_metadata: @@ -307,7 +265,7 @@ class StepContextGenerator(BaseContextGenerator): candidate_text = f"Candidate {i}:\n" candidate_text += f"Condition: {condition}\n" candidate_text += f"Experience: {experience}\n" - candidate_text += f"Tags: {', '.join(tags)}\n" + candidate_text += f"Tags: {', '.join(tags) if tags else 'None'}\n" candidate_text += f"Confidence: {confidence}\n" formatted_candidates.append(candidate_text) @@ -349,27 +307,31 @@ class StepContextGenerator(BaseContextGenerator): exp_text = f"Experience {i} ({experience_type}):\n" exp_text += f"When to use: {condition}\n" exp_text += f"Experience: {experience_content}\n" - exp_text += f"Tags: {', '.join(tags)}" + exp_text += f"Tags: {', '.join(tags) if tags else 'None'}" formatted_experiences.append(exp_text) return "\n\n---\n\n".join(formatted_experiences) - def _create_context(self, experiences: List[VectorStoreNode]) -> str: + def _create_simple_context(self, experiences: List[VectorStoreNode]) -> str: """Create simple context when LLM generation fails""" if not experiences: return "" - context = "Here are some relevant experiences that might help:\n\n" + context = "Previous Experience\n" - for i, exp in enumerate(experiences, 1): + for exp in experiences: condition = exp.content experience_content = exp.metadata.get("experience", "") - context += f"{i}. **When**: {condition}\n" - context += f" **Experience**: {experience_content}\n\n" + if not experience_content: + continue - return context + context += f"- {condition} {experience_content}\n" + + context += "Please consider the helpful parts from these in answering the question, to make the response more comprehensive and substantial." + + return context.strip() def _parse_json_response(self, response: str, key: str) -> str: """Parse JSON response to extract specific key""" @@ -391,4 +353,4 @@ class StepContextGenerator(BaseContextGenerator): except json.JSONDecodeError: logger.warning(f"Failed to parse JSON response for key '{key}'") - return "" + return "" \ No newline at end of file diff --git a/experiencemaker/module/summarizer/step_summarizer.py b/experiencemaker/module/summarizer/step_summarizer.py index 9caa2e3d..9ca01e46 100644 --- a/experiencemaker/module/summarizer/step_summarizer.py +++ b/experiencemaker/module/summarizer/step_summarizer.py @@ -2,60 +2,77 @@ import json import re import uuid from datetime import datetime +from pathlib import Path from typing import List, Dict, Any, Optional, Tuple from loguru import logger -from pydantic import Field, model_validator +from pydantic import Field from experiencemaker.enumeration.role import Role -from experiencemaker.module.summarizer.base_summarizer import BaseSummarizer -from experiencemaker.schema.trajectory import Trajectory, Sample, SummaryMessage, Message -from experiencemaker.schema.vector_store_node import VectorStoreNode -from experiencemaker.storage.es_vector_store import EsVectorStore -from experiencemaker.storage.file_vector_store import FileVectorStore +from experiencemaker.module.prompt.prompt_mixin import PromptMixin +from experiencemaker.module.summarizer.base_summarizer import BaseSummarizer, SUMMARIZER_REGISTRY +from experiencemaker.schema.experience import Experience +from experiencemaker.schema.trajectory import Trajectory, Message -class StepSummarizer(BaseSummarizer): +@SUMMARIZER_REGISTRY.register("step") +class StepSummarizer(BaseSummarizer, PromptMixin): """ Step-level experience extractor that focuses on extracting reusable experiences from individual steps or step sequences in trajectories """ - # Vector Store 配置 - vector_store_type: str = Field(default="file_vector_store") - vector_store_hosts: str | List[str] = Field(default="http://localhost:9200") - vector_store_index_name: str = Field(default="step_experience_store") - store_dir: str = Field(default="./step_experiences/") + # Feature switches - can be configured via startup parameters + enable_step_segmentation: bool = Field(default=False, description="Enable trajectory segmentation into steps") + enable_similarity_search: bool = Field(default=False, description="Enable similarity search for comparison") + enable_experience_validation: bool = Field(default=True, description="Enable experience validation") - # 功能开关 - enable_step_segmentation: bool = Field(default=False) - enable_similarity_search: bool = Field(default=False) - enable_experience_validation: bool = Field(default=True) + # LLM retries + max_retries: int = Field(default=3, description="Maximum retries for LLM calls") - # llm retries - max_retries: int = Field(default=3) + # Prompt configuration + prompt_file_path: Path = Field(default=Path(__file__).parent / "step_summarizer_prompt.yaml") - @model_validator(mode="after") - def init_vector_store(self): - """initialize""" - if self.vector_store_type == "file_vector_store": - self.vector_store = FileVectorStore( - embedding_model=self.embedding_model, - index_name=self.vector_store_index_name, - store_dir=self.store_dir - ) - elif self.vector_store_type == "es_vector_store": - self.vector_store = EsVectorStore( - embedding_model=self.embedding_model, - index_name=self.vector_store_index_name, - hosts=self.vector_store_hosts + def _extract_experiences(self, trajectories: List[Trajectory], workspace_id: str = None, + **kwargs) -> List[Experience]: + """Extract step-level experiences from trajectories (implements base class method)""" + logger.info(f"Starting step-level experience extraction pipeline for {len(trajectories)} trajectories") + + all_experiences = [] + + # Classify trajectories based on trajectory.done + success_trajectories = [traj for traj in trajectories if traj.done] + failure_trajectories = [traj for traj in trajectories if not traj.done] + + # Process success and failure samples separately + if success_trajectories: + success_experiences = self._extract_step_experiences_from_success(success_trajectories, workspace_id, + **kwargs) + all_experiences.extend(success_experiences) + + if failure_trajectories: + failure_experiences = self._extract_step_experiences_from_failure(failure_trajectories, workspace_id, + **kwargs) + all_experiences.extend(failure_experiences) + + # Comparative analysis (if similarity search is enabled) + if success_trajectories and failure_trajectories and self.enable_similarity_search: + comparative_experiences = self._extract_step_experiences_from_comparison( + success_trajectories, failure_trajectories, workspace_id, **kwargs ) + all_experiences.extend(comparative_experiences) + + # Validate experiences + if self.enable_experience_validation: + validated_experiences = self._validate_experiences(all_experiences, **kwargs) else: - raise ValueError(f"Unknown vector store type: {self.vector_store_type}") + validated_experiences = all_experiences - return self + logger.info(f"Extracted {len(validated_experiences)} validated step experiences") + return validated_experiences - def extract_step_experiences_from_success(self, trajectories: List[Trajectory], **kwargs) -> List[SummaryMessage]: + def _extract_step_experiences_from_success(self, trajectories: List[Trajectory], workspace_id: str, **kwargs) -> \ + List[Experience]: """Extract step-level experiences from successful samples""" logger.info(f"Extracting step experiences from {len(trajectories)} successful trajectories") @@ -65,16 +82,36 @@ class StepSummarizer(BaseSummarizer): for step_seq in step_sequences: try: - prompt = self.prompt_handler.success_step_experience_prompt.format( + step_content_collector = [] + for step in step_seq: + step_index = len(step_content_collector) + + if step.role is Role.ASSISTANT: + line = f"### step.{step_index} role={step.role.value} content=\n{step.content}\n" + if hasattr(step, 'reasoning_content') and step.reasoning_content: + line += f"{step.reasoning_content}\n" + if hasattr(step, 'tool_calls') and step.tool_calls: + for tool_call in step.tool_calls: + line += f" - tool call={tool_call.name}\n params={tool_call.arguments}\n" + step_content_collector.append(line) + elif step.role is Role.USER: + line = f"### step.{step_index} role={step.role.value} content=\n{step.content}\n" + step_content_collector.append(line) + elif step.role is Role.TOOL: + line = f"### step.{step_index} role={step.role.value} tool call result=\n{step.content}\n" + step_content_collector.append(line) + + prompt = self.prompt_format( + prompt_name="success_step_experience_prompt", query=trajectory.query, - step_sequence=self._format_step_sequence(step_seq), + step_sequence="\n".join(step_content_collector).strip(), context=self._get_trajectory_context(trajectory, step_seq), outcome="successful" ) - experience = self._extract_with_llm(prompt, "success") - if experience: - all_experiences.extend(experience) + experiences = self._extract_with_llm(prompt, "success", workspace_id) + if experiences: + all_experiences.extend(experiences) except Exception as e: logger.error(f"Error extracting success experience: {e}") @@ -82,7 +119,8 @@ class StepSummarizer(BaseSummarizer): return all_experiences - def extract_step_experiences_from_failure(self, trajectories: List[Trajectory], **kwargs) -> List[SummaryMessage]: + def _extract_step_experiences_from_failure(self, trajectories: List[Trajectory], workspace_id: str, **kwargs) -> \ + List[Experience]: """Extract step-level experiences from failed samples""" logger.info(f"Extracting step experiences from {len(trajectories)} failed trajectories") @@ -92,16 +130,36 @@ class StepSummarizer(BaseSummarizer): for step_seq in step_sequences: try: - prompt = self.prompt_handler.failure_step_experience_prompt.format( + step_content_collector = [] + for step in step_seq: + step_index = len(step_content_collector) + + if step.role is Role.ASSISTANT: + line = f"### step.{step_index} role={step.role.value} content=\n{step.content}\n" + if hasattr(step, 'reasoning_content') and step.reasoning_content: + line += f"{step.reasoning_content}\n" + if hasattr(step, 'tool_calls') and step.tool_calls: + for tool_call in step.tool_calls: + line += f" - tool call={tool_call.name}\n params={tool_call.arguments}\n" + step_content_collector.append(line) + elif step.role is Role.USER: + line = f"### step.{step_index} role={step.role.value} content=\n{step.content}\n" + step_content_collector.append(line) + elif step.role is Role.TOOL: + line = f"### step.{step_index} role={step.role.value} tool call result=\n{step.content}\n" + step_content_collector.append(line) + + prompt = self.prompt_format( + prompt_name="failure_step_experience_prompt", query=trajectory.query, - step_sequence=self._format_step_sequence(step_seq), + step_sequence="\n".join(step_content_collector).strip(), context=self._get_trajectory_context(trajectory, step_seq), outcome="failed" ) - experience = self._extract_with_llm(prompt, "failure") - if experience: - all_experiences.extend(experience) + experiences = self._extract_with_llm(prompt, "failure", workspace_id) + if experiences: + all_experiences.extend(experiences) except Exception as e: logger.error(f"Error extracting failure experience: {e}") @@ -109,10 +167,11 @@ class StepSummarizer(BaseSummarizer): return all_experiences - def extract_step_experiences_from_comparison(self, - success_trajectories: List[Trajectory], - failure_trajectories: List[Trajectory], - **kwargs) -> List[SummaryMessage]: + def _extract_step_experiences_from_comparison(self, + success_trajectories: List[Trajectory], + failure_trajectories: List[Trajectory], + workspace_id: str, + **kwargs) -> List[Experience]: """Extract step-level experiences from comparative samples""" logger.info(f"Extracting comparative step experiences from {len(success_trajectories)} success " f"and {len(failure_trajectories)} failure trajectories") @@ -124,15 +183,16 @@ class StepSummarizer(BaseSummarizer): for success_steps, failure_steps, similarity_score in similar_step_pairs: try: - prompt = self.prompt_handler.comparative_step_experience_prompt.format( + prompt = self.prompt_format( + prompt_name="comparative_step_experience_prompt", success_steps=self._format_step_sequence(success_steps), failure_steps=self._format_step_sequence(failure_steps), similarity_score=similarity_score ) - experience = self._extract_with_llm(prompt, "comparative") - if experience: - all_experiences.extend(experience) + experiences = self._extract_with_llm(prompt, "comparative", workspace_id) + if experiences: + all_experiences.extend(experiences) except Exception as e: logger.error(f"Error extracting comparative experience: {e}") @@ -140,34 +200,7 @@ class StepSummarizer(BaseSummarizer): return all_experiences - def extract_step_experiences_general(self, trajectories: List[Trajectory], **kwargs) -> List[SummaryMessage]: - """Extract general step experiences when no labels are provided""" - logger.info(f"Extracting general step experiences from {len(trajectories)} trajectories") - - all_experiences = [] - - for trajectory in trajectories: - step_sequences = self._segment_trajectory_into_steps(trajectory) - - for step_seq in step_sequences: - try: - prompt = self.prompt_handler.general_step_experience_prompt.format( - query=trajectory.query, - step_sequence=self._format_step_sequence(step_seq), - context=self._get_trajectory_context(trajectory, step_seq) - ) - - experience = self._extract_with_llm(prompt, "general") - if experience: - all_experiences.extend(experience) - - except Exception as e: - logger.error(f"Error extracting general experience: {e}") - continue - - return all_experiences - - def validate_experiences(self, experiences: List[SummaryMessage], **kwargs) -> List[SummaryMessage]: + def _validate_experiences(self, experiences: List[Experience], **kwargs) -> List[Experience]: """Validate the quality and validity of extracted experiences""" if not self.enable_experience_validation: return experiences @@ -181,12 +214,6 @@ class StepSummarizer(BaseSummarizer): validation_result = self._validate_single_experience(experience) if validation_result["is_valid"]: - # Add validation info to metadata - experience.metadata.update({ - "validation_score": validation_result["score"], - "validation_feedback": validation_result["feedback"], - "validated_at": datetime.now().isoformat() - }) validated_experiences.append(experience) else: logger.warning(f"Experience validation failed: {validation_result['reason']}") @@ -198,73 +225,6 @@ class StepSummarizer(BaseSummarizer): logger.info(f"Validated {len(validated_experiences)} out of {len(experiences)} experiences") return validated_experiences - def store_experiences(self, experiences: List[SummaryMessage], **kwargs): - """Store experiences into vector storage""" - if not experiences: - logger.warning("No experiences to store") - return - - # Deduplication - unique_experiences = self._deduplicate_experiences(experiences) - logger.info(f"Storing {len(unique_experiences)} unique experiences (deduplicated from {len(experiences)})") - - # Convert to storage nodes - nodes = [] - for exp in unique_experiences: - node = VectorStoreNode( - content=exp.content, - metadata={ - **exp.metadata, - "stored_at": datetime.now().isoformat(), - "experience_type": "step_level" - } - ) - nodes.append(node) - - # Store to vector database - refresh_index = kwargs.get("refresh_index", True) - self.vector_store.insert(nodes, refresh_index=refresh_index) - logger.info(f"Successfully stored {len(nodes)} step experiences") - - def execute(self, trajectories: List[Trajectory], **kwargs) -> List[Sample]: - """Execute complete step-level experience extraction pipeline""" - logger.info(f"Starting step-level experience extraction pipeline for {len(trajectories)} trajectories") - - all_experiences = [] - - # Classify trajectories based on trajectory.done - success_trajectories = [traj for traj in trajectories if traj.done] - failure_trajectories = [traj for traj in trajectories if not traj.done] - - # Process success and failure samples separately - if success_trajectories: - success_experiences = self.extract_step_experiences_from_success(success_trajectories, **kwargs) - all_experiences.extend(success_experiences) - - if failure_trajectories: - failure_experiences = self.extract_step_experiences_from_failure(failure_trajectories, **kwargs) - all_experiences.extend(failure_experiences) - - # Comparative analysis (if similarity search is enabled) - if success_trajectories and failure_trajectories and self.enable_similarity_search: - comparative_experiences = self.extract_step_experiences_from_comparison( - success_trajectories, failure_trajectories, **kwargs - ) - all_experiences.extend(comparative_experiences) - - # Validate experiences - if self.enable_experience_validation: - validated_experiences = self.validate_experiences(all_experiences, **kwargs) - else: - validated_experiences = all_experiences - - # Store experiences - if validated_experiences: - self.store_experiences(validated_experiences, **kwargs) - - # Construct return result - return [Sample(steps=validated_experiences)] - # ========== Helper Methods ========== def _segment_trajectory_into_steps(self, trajectory: Trajectory) -> List[List[Message]]: @@ -277,7 +237,8 @@ class StepSummarizer(BaseSummarizer): # Use LLM for segmentation trajectory_content = self._format_trajectory_content(trajectory) - prompt = self.prompt_handler.step_segmentation_prompt.format( + prompt = self.prompt_format( + prompt_name="step_segmentation_prompt", query=trajectory.query, trajectory_content=trajectory_content, total_steps=len(trajectory.steps) @@ -413,9 +374,9 @@ class StepSummarizer(BaseSummarizer): success_texts = [self._format_step_sequence(seq) for seq in success_step_sequences] failure_texts = [self._format_step_sequence(seq) for seq in failure_step_sequences] - # Get embeddings - success_embeddings = self.embedding_model.get_embeddings(success_texts) - failure_embeddings = self.embedding_model.get_embeddings(failure_texts) + # Get embeddings using embedding model + success_embeddings = self.vector_store.embedding_model.get_embeddings(success_texts) + failure_embeddings = self.vector_store.embedding_model.get_embeddings(failure_texts) # Calculate similarity and find most similar pairs for i, s_emb in enumerate(success_embeddings): @@ -457,16 +418,29 @@ class StepSummarizer(BaseSummarizer): except Exception as e: logger.error(f"Error calculating cosine similarity: {e}") return 0.0 - import json - def _extract_with_llm(self, prompt: str, experience_type: str) -> List[SummaryMessage]: + def _extract_with_llm(self, prompt: str, experience_type: str, workspace_id: str) -> List[Experience]: + """Extract experiences using LLM with JSON parsing - can return multiple experiences""" for attempt in range(self.max_retries): try: response = self.llm.chat([Message(role=Role.USER, content=prompt)]) - experiences = self._parse_experience_response(response.content, experience_type) - if experiences: + # Parse JSON response to extract experiences + experiences_data = self._parse_json_experience_response(response.content) + + if experiences_data: + experiences = [] + for exp_data in experiences_data: + experience = Experience( + experience_workspace_id=workspace_id, + experience_desc=exp_data.get("condition", exp_data.get("when_to_use", "")), + experience_content=exp_data.get("experience", ""), + metadata = exp_data + ) + experiences.append(experience) return experiences + else: + logger.warning(f"Experience extraction failed: no valid JSON experience found in response") except Exception as e: logger.warning(f"Attempt {attempt + 1} failed for experience extraction: {e}") @@ -474,72 +448,58 @@ class StepSummarizer(BaseSummarizer): logger.error(f"Failed to extract experience after {self.max_retries} attempts") return [] - def _parse_experience_response(self, response: str, experience_type: str) -> List[SummaryMessage]: - """解析经验抽取响应""" - experiences = [] - + def _parse_json_experience_response(self, response: str) -> List[dict]: + """Parse JSON experience response - handles both single objects and arrays""" try: - # 尝试提取JSON格式的经验 + # Try to extract JSON format json_pattern = r'```json\s*([\s\S]*?)\s*```' json_blocks = re.findall(json_pattern, response) - for block in json_blocks: - try: - parsed = json.loads(block) - if isinstance(parsed, list): - for exp_data in parsed: - experience = self._create_experience_message(exp_data, experience_type) - if experience: - experiences.append(experience) - else: - experience = self._create_experience_message(parsed, experience_type) - if experience: - experiences.append(experience) - except json.JSONDecodeError: - continue + if json_blocks: + parsed = json.loads(json_blocks[0]) - except Exception as e: - logger.error(f"Error parsing experience response: {e}") + # Handle array of experiences + if isinstance(parsed, list): + valid_experiences = [] + for exp_data in parsed: + if isinstance(exp_data, dict) and ( + ("condition" in exp_data and "experience" in exp_data) or + ("when_to_use" in exp_data and "experience" in exp_data) + ): + valid_experiences.append(exp_data) + return valid_experiences - return experiences + # Handle single experience object + elif isinstance(parsed, dict) and ( + ("condition" in parsed and "experience" in parsed) or + ("when_to_use" in parsed and "experience" in parsed) + ): + return [parsed] - def _create_experience_message(self, exp_data: Dict[str, Any], experience_type: str) -> Optional[SummaryMessage]: - """创建经验消息对象""" + # Fallback: try to parse the entire response as JSON + parsed = json.loads(response) + if isinstance(parsed, list): + return parsed + elif isinstance(parsed, dict): + return [parsed] + + except json.JSONDecodeError as e: + logger.warning(f"Failed to parse JSON experience response: {e}") + + return [] + + def _validate_single_experience(self, experience: Experience) -> Dict[str, Any]: + """Validate single experience""" try: - condition = exp_data.get("when_to_use", exp_data.get("condition", "")) - experience_content = exp_data.get("experience", exp_data.get("tip_content", exp_data.get("tips", ""))) - - if not condition or not experience_content: - return None - - metadata = { - "experience": experience_content, - "experience_type": experience_type, - "tags": exp_data.get("tags", []), - "confidence": exp_data.get("confidence", 0.5), - "extracted_at": datetime.now().isoformat(), - "experience_id": str(uuid.uuid4()) - } - - return SummaryMessage(content=condition, metadata=metadata) - - except Exception as e: - logger.error(f"Error creating experience message: {e}") - return None - - def _validate_single_experience(self, experience: SummaryMessage) -> Dict[str, Any]: - """验证单个经验的有效性""" - try: - prompt = self.prompt_handler.experience_validation_prompt.format( - condition=experience.content, - experience_content=experience.metadata.get("experience", ""), - experience_type=experience.metadata.get("experience_type", ""), - tags=experience.metadata.get("tags", []) + prompt = self.prompt_format( + prompt_name="experience_validation_prompt", + condition=experience.experience_desc, + experience_content=experience.experience_content, ) response = self.llm.chat([Message(role=Role.USER, content=prompt)]) - # 解析验证结果 + # Parse validation result is_valid = "valid" in response.content.lower() and "invalid" not in response.content.lower() score_match = re.search(r'score[:\s]*([0-9.]+)', response.content.lower()) score = float(score_match.group(1)) if score_match else 0.5 @@ -553,29 +513,4 @@ class StepSummarizer(BaseSummarizer): except Exception as e: logger.error(f"Error validating experience: {e}") - return {"is_valid": False, "score": 0.0, "feedback": "", "reason": str(e)} - - def _deduplicate_experiences(self, experiences: List[SummaryMessage]) -> List[SummaryMessage]: - unique_experiences = [] - seen_contents = set() - - for exp in experiences: - content_hash = hash(exp.content) - - if content_hash not in seen_contents: - seen_contents.add(content_hash) - unique_experiences.append(exp) - - return unique_experiences - - def extract_samples(self, trajectories: List[Trajectory], **kwargs) -> List[Sample]: - experiences = self.execute(trajectories, **kwargs) - return [Sample(steps=experiences)] if experiences else [] - - def insert_into_vector_store(self, samples: List[Sample], **kwargs): - all_experiences = [] - for sample in samples: - all_experiences.extend(sample.steps) - - if all_experiences: - self.store_experiences(all_experiences, **kwargs) + return {"is_valid": False, "score": 0.0, "feedback": "", "reason": str(e)} \ No newline at end of file