diff --git a/kag/common/tools/search_api/impl/openspg_search_api.py b/kag/common/tools/search_api/impl/openspg_search_api.py index c7a4a5455..37eefb665 100644 --- a/kag/common/tools/search_api/impl/openspg_search_api.py +++ b/kag/common/tools/search_api/impl/openspg_search_api.py @@ -1,9 +1,12 @@ +import logging from typing import List from kag.common.conf import KAGConstants, KAGConfigAccessor from kag.common.tools.search_api.search_api_abc import SearchApiABC from knext.search.client import SearchClient +logger = logging.getLogger(__name__) + @SearchApiABC.register("openspg_search_api", as_default=True) class OpenSPGSearchAPI(SearchApiABC): @@ -25,12 +28,15 @@ def search_text( self, query_string, label_constraints=None, topk=10, params=None ) -> List: if self.sc: - return self.sc.search_text( - query_string=query_string, - label_constraints=label_constraints, - topk=topk, - params=params, - ) + try: + return self.sc.search_text( + query_string=query_string, + label_constraints=label_constraints, + topk=topk, + params=params, + ) + except Exception as e: + logger.debug(f"search_vector error {e}", exc_info=True) return [] def search_vector( @@ -38,14 +44,18 @@ def search_vector( ) -> List: if self.sc is None: return [] - res = self.sc.search_vector( - label=label, - property_key=property_key, - query_vector=query_vector, - topk=topk, - ef_search=ef_search, - params=params, - ) + try: + res = self.sc.search_vector( + label=label, + property_key=property_key, + query_vector=query_vector, + topk=topk, + ef_search=ef_search, + params=params, + ) + except Exception as e: + logger.debug(f"search_vector error {e}", exc_info=True) + res = [] if res is None: return [] return res diff --git a/kag/interface/common/llm_client.py b/kag/interface/common/llm_client.py index b9738b0e8..22390246b 100644 --- a/kag/interface/common/llm_client.py +++ b/kag/interface/common/llm_client.py @@ -326,6 +326,8 @@ def invoke( tools = kwargs.get("tools", None) if tools: with_json_parse = False + context = kwargs.get("context", None) + segment_name = kwargs.get("segment_name", None) try: self.sync_limiter.acquire() response = ( @@ -335,6 +337,11 @@ def invoke( ) if tools: return response + if context and segment_name and hasattr(context, "add_kwargs"): + try: + context.add_kwargs(segment_name, response) + except Exception as e: + logger.warning(f"Failed to add kwargs to context: {e}") result = prompt_op.parse_response(response, model=self.model, **variables) logger.debug(f"Result: {result}") return result diff --git a/kag/interface/solver/context.py b/kag/interface/solver/context.py index d0eb617d6..12793af24 100644 --- a/kag/interface/solver/context.py +++ b/kag/interface/solver/context.py @@ -14,6 +14,23 @@ def __init__(self): self.variables_graph = KgGraph() self.kwargs = {} + def add_kwargs(self, key, value): + """ + Adds a key-value pair to the context's keyword arguments dictionary. + + This method stores arbitrary key-value data in the context that can be + accessed by other components during pipeline execution. + + Args: + key: The string key to store the value under + value: The value to be stored, can be of any type + + Example: + context.add_kwargs("temperature", 0.7) + context.add_kwargs("max_tokens", 1000) + """ + self.kwargs[key] = value + def add_task(self, task: Task): """Adds a task to the context. diff --git a/kag/solver/pipeline/kag_static_pipeline.py b/kag/solver/pipeline/kag_static_pipeline.py index 26f7a323e..4ba7ce199 100644 --- a/kag/solver/pipeline/kag_static_pipeline.py +++ b/kag/solver/pipeline/kag_static_pipeline.py @@ -129,21 +129,26 @@ async def ainvoke(self, query, **kwargs): while True: context: Context = Context() tasks = await self.planning(query, context, **kwargs) - - for task in tasks: - context.add_task(task) - - for task_group in context.gen_task(group=True): - await asyncio.gather( - *[ - asyncio.create_task( - self.execute_task(query, task, context, **kwargs) - ) - for task in task_group - ] + if not tasks: + think_response = context.kwargs.get("thinker", "") + answer = await self.generator.ainvoke( + query + "\n" + think_response, context, **kwargs ) + else: + for task in tasks: + context.add_task(task) + + for task_group in context.gen_task(group=True): + await asyncio.gather( + *[ + asyncio.create_task( + self.execute_task(query, task, context, **kwargs) + ) + for task in task_group + ] + ) - answer = await self.generator.ainvoke(query, context, **kwargs) + answer = await self.generator.ainvoke(query, context, **kwargs) from kag.common.utils import red, green, reset task_info = [] diff --git a/kag/solver/reporter/open_spg_reporter.py b/kag/solver/reporter/open_spg_reporter.py index 8d0d9c95f..6877313a6 100644 --- a/kag/solver/reporter/open_spg_reporter.py +++ b/kag/solver/reporter/open_spg_reporter.py @@ -530,8 +530,6 @@ def add_report_line(self, segment, tag_name, content, status, **kwargs): } def do_report(self): - if not self.client: - return content, status_enum, metrics = self.generate_report_data() request = TaskStreamRequest( @@ -539,9 +537,12 @@ def do_report(self): ) # logging.info(f"do_report:{request}") try: - ret = self.client.reasoner_dialog_report_completions_post( - task_stream_request=request - ) + if self.client: + ret = self.client.reasoner_dialog_report_completions_post( + task_stream_request=request + ) + else: + ret = {} if self.last_report is None: logger.info(f"begin do_report: {request} ret={ret}") self.last_report = request diff --git a/requirements.txt b/requirements.txt index 44ed0d4b9..27a1c4a98 100644 --- a/requirements.txt +++ b/requirements.txt @@ -45,6 +45,7 @@ json_repair gitpython docstring_parser aiolimiter +pyarrow==19.0.1 pyodps==0.12.2 aliyun-log-python-sdk==0.8.8 pyvis