ref: caching, session

AaLexUser · AaLexUser · commit 588a8eb00b6d · 2025-07-15T22:22:42.000+03:00
diff --git a/fedotllm/agents/automl/automl.py b/fedotllm/agents/automl/automl.py
@@ -43,7 +43,7 @@ def __init__(
         self, config: AppConfig, dataset_path: str | Path, workspace: str | Path
     ):
         self.config = config
-        self.inference = AIInference(config.llm)
+        self.inference = AIInference(config.llm, config.session_id)
         self.dataset = Dataset.from_path(dataset_path)
         self.workspace = Path(workspace)
 
diff --git a/fedotllm/agents/automl/nodes.py b/fedotllm/agents/automl/nodes.py
@@ -216,13 +216,19 @@ def if_bug(state: AutoMLAgentState, app_config: AppConfig):
 def fix_solution(state: AutoMLAgentState, inference: AIInference, dataset: Dataset):
     logger.info("Running fix solution")
 
+    stdout_lines = state["observation"].stdout.splitlines()
+    if len(stdout_lines) > 20:
+        stdout = "\n".join(stdout_lines[:10] + ["..."] + stdout_lines[-10:])
+    else:
+        stdout = state["observation"].stdout
+
     fix_prompt = prompts.automl.fix_solution_prompt(
         reflection=state["reflection"],
         dataset_path=str(dataset.path.absolute()),
         code_recent_solution=state["raw_code"],
         msg=state["observation"].msg,
         stderr=state["observation"].stderr,
-        stdout=state["observation"].stdout,
+        stdout=stdout,
     )
 
     fixed_solution = inference.query(fix_prompt)
diff --git a/fedotllm/agents/researcher/researcher.py b/fedotllm/agents/researcher/researcher.py
@@ -14,6 +14,7 @@
     rewrite_question,
 )
 from fedotllm.agents.researcher.state import ResearcherAgentState
+from fedotllm.configs.schema import AppConfig
 from fedotllm.llm import AIInference, LiteLLMEmbeddings
 
 RETRIEVE = "retrieve"
@@ -24,9 +25,9 @@
 
 
 class ResearcherAgent(Agent):
-    def __init__(self, inference: AIInference, embeddings: LiteLLMEmbeddings):
-        self.inference = inference
-        self.embeddings = embeddings
+    def __init__(self, config: AppConfig):
+        self.inference = AIInference(config.llm, config.session_id)
+        self.embeddings = LiteLLMEmbeddings(config.embeddings)
 
     def create_graph(self):
         workflow = StateGraph(ResearcherAgentState)
diff --git a/fedotllm/agents/supervisor.py b/fedotllm/agents/supervisor.py
@@ -9,6 +9,7 @@
 from pydantic import BaseModel, Field
 
 from fedotllm.agents.base import Agent, FedotLLMAgentState
+from fedotllm.configs.schema import AppConfig
 from fedotllm.llm import AIInference
 from fedotllm.prompts.supervisor import choose_next_prompt
 
@@ -26,11 +27,11 @@ class SupervisorState(FedotLLMAgentState):
 class SupervisorAgent(Agent):
     def __init__(
         self,
-        inference: AIInference,
+        config: AppConfig,
         automl_agent: Runnable,
         researcher_agent: Runnable,
     ):
-        self.inference = inference
+        self.inference = AIInference(config.llm, config.session_id)
         self.researcher_agent = researcher_agent
         self.automl_agent = automl_agent
 
@@ -50,7 +51,7 @@ def finish_execution(state: SupervisorState):
         workflow.add_edge("researcher", "choose_next")
         workflow.add_edge("automl", "finish")
         workflow.add_edge("finish", END)
-        return workflow.compile().with_config(config={"run_name": "SupervisorAgent"})
+        return workflow.compile().with_config(run_name=SupervisorAgent)
 
 
 class ChooseNext(BaseModel):
@@ -72,6 +73,7 @@ def router_node(
     """
 
     messages = convert_to_openai_messages(state["messages"])
+    messages = [messages] if isinstance(messages, dict) else messages
     messages.append({"role": "user", "content": choose_next_prompt()})
 
     response = inference.query(messages)
diff --git a/fedotllm/configs/default.yaml b/fedotllm/configs/default.yaml
@@ -12,6 +12,8 @@ llm:
   model_name: gemini-2.0-flash
   base_url: https://generativelanguage.googleapis.com/v1beta/openai/
   api_key: ${oc.env:FEDOTLLM_LLM_API_KEY}
+  caching:
+    enabled: true
   extra_headers:
     X-Title: FEDOT.LLM
 embeddings:
diff --git a/fedotllm/configs/loader.py b/fedotllm/configs/loader.py
@@ -1,6 +1,5 @@
 import logging
 import re
-from importlib.resources import files
 from pathlib import Path
 from typing import List, Optional, Type, TypeVar
 
@@ -9,6 +8,7 @@
 from pydantic import BaseModel
 
 from fedotllm.configs.schema import AppConfig
+from fedotllm.constants import PACKAGE_PATH
 
 
 def _get_default_config_path(
@@ -27,7 +27,7 @@ def _get_default_config_path(
         ValueError: If the config file is not found.
     """
     try:
-        config_path = Path(files("fedotllm") / "configs" / f"{presets}.yaml")
+        config_path = Path(PACKAGE_PATH) / "configs" / f"{presets}.yaml"
 
         if not config_path.exists():
             raise ValueError(
diff --git a/fedotllm/configs/ollama.yaml b/fedotllm/configs/ollama.yaml
@@ -0,0 +1,5 @@
+llm:
+  provider: ollama
+  model_name: "devstral:latest"
+  base_url:
+  api_key: "ollama"
diff --git a/fedotllm/configs/schema.py b/fedotllm/configs/schema.py
@@ -1,7 +1,10 @@
+from pathlib import Path
 from typing import Any, Dict, Optional
 
 from pydantic import BaseModel, Field
 
+from fedotllm.constants import PACKAGE_PATH
+
 
 class TemplatesConfig(BaseModel):
     code: str
@@ -16,11 +19,17 @@ class AutoMLConfig(BaseModel):
     predictor_init_kwargs: dict = Field(default_factory=dict)
 
 
+class CachingConfig(BaseModel):
+    enabled: bool = True
+    dir_path: str = Field(default=str(Path(PACKAGE_PATH) / "cache"))
+
+
 class LLMConfig(BaseModel):
     provider: str = "openai"
     model_name: str = "gpt-4o"
     base_url: Optional[str] = None
     api_key: Optional[str] = None
+    caching: CachingConfig = Field(default_factory=CachingConfig)
     extra_headers: Dict[str, Any] = {}
     completion_params: Dict[str, Any] = {}
 
@@ -45,3 +54,4 @@ class AppConfig(BaseModel):
     embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
     langfuse: LangfuseConfig = Field(default_factory=LangfuseConfig)
     automl: AutoMLConfig = Field(default_factory=AutoMLConfig)
+    session_id: Optional[str] = Field(default=None)
diff --git a/fedotllm/constants.py b/fedotllm/constants.py
@@ -1,3 +1,8 @@
+from importlib.resources import files
+from pathlib import Path
+
+PACKAGE_PATH = str(Path(files("fedotllm")))
+
 # File formats
 CSV_SUFFIXES = [".csv"]
 PARQUET_SUFFIXES = [".parquet", ".pq"]
diff --git a/fedotllm/data.py b/fedotllm/data.py
@@ -126,7 +126,7 @@ def dataset_eda(self):
         train_split = self.get_train_split()
         df = train_split.data
         eda = ""
-        if df.shape[1] <= 10:
+        if df.shape[1] <= 20:
             eda += "\n===== 1. BASIC INFO =====\n"
 
             buf = io.StringIO()
diff --git a/fedotllm/llm.py b/fedotllm/llm.py
@@ -1,9 +1,10 @@
 import os
-from typing import Any, Dict, List, Type, TypeVar
+from typing import Any, Dict, List, Optional, Type, TypeVar
 
 import litellm
 import tiktoken
-from pydantic import BaseModel
+from litellm.caching.caching import Cache, LiteLLMCacheType
+from pydantic import BaseModel, ValidationError
 from tenacity import retry, stop_after_attempt, wait_exponential
 
 from fedotllm import prompts
@@ -24,7 +25,7 @@
 
 
 class AIInference:
-    def __init__(self, config: LLMConfig):
+    def __init__(self, config: LLMConfig, session_id: Optional[str] = None):
         self.config = config
 
         if not self.config.api_key:
@@ -37,9 +38,15 @@ def __init__(self, config: LLMConfig):
             "api_key": self.config.api_key,
             "base_url": self.config.base_url,
             "extra_headers": self.config.extra_headers,
+            "metadata": {"session_id": session_id},
             **self.config.completion_params,
         }
 
+        if config.caching.enabled:
+            litellm.cache = Cache(
+                type=LiteLLMCacheType.DISK, disk_cache_dir=config.caching.dir_path
+            )
+
     @retry(
         stop=stop_after_attempt(5),
         wait=wait_exponential(multiplier=1, min=4, max=10),
@@ -49,7 +56,13 @@ def create(self, messages: str, response_model: Type[T]) -> T:
         messages = f"{messages}\n{prompts.utils.structured_response(response_model)}"
         response = self.query(messages)
         json_obj = parse_json(response) if response else None
-        return response_model.model_validate(json_obj)
+        try:
+            return response_model.model_validate(json_obj)
+        except ValidationError as exc_info:
+            messages = f"{prompts.utils.fix_structured_response(json_obj, str(exc_info), response_model)}"
+            response = self.query(messages)
+            json_obj = parse_json(response) if response else None
+            return response_model.model_validate(json_obj)
 
     @retry(
         stop=stop_after_attempt(5),
diff --git a/fedotllm/log.py b/fedotllm/log.py
@@ -1,7 +1,7 @@
 import logging
 
 logger = logging.getLogger("FEDOTLLM")
-logger.setLevel(logging.DEBUG)  # Changed from INFO to DEBUG to allow debug messages
+logger.setLevel(logging.DEBUG)
 
 file_handler = logging.FileHandler("fedotllm.log", mode="w")
 file_handler.setLevel(logging.DEBUG)
diff --git a/fedotllm/main.py b/fedotllm/main.py
@@ -51,16 +51,17 @@ def __init__(
         self.inference = AIInference(self.config.llm)
         self.embeddings = LiteLLMEmbeddings(self.config.embeddings)
         self.handlers = handlers if handlers is not None else []
+        self.config.session_id = self.config.session_id or pd.Timestamp.now().strftime(
+            "%Y%m%d_%H%M%S"
+        )
         self.workspace = workspace
 
     async def ainvoke(self, message: str):
         logger.info(
             f"FedotAI ainvoke called. Input message (first 100 chars): '{message[:100]}...'"
         )
         if not self.workspace:
-            self.workspace = Path(
-                f"fedotllm-output-{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}"
-            )
+            self.workspace = Path(f"fedotllm-output-{self.config.session_id}")
             logger.info(f"Workspace for ainvoke created at: {self.workspace}")
 
         translator_agent = TranslatorAgent(inference=self.inference)
@@ -74,12 +75,10 @@ async def ainvoke(self, message: str):
         automl_agent = AutoMLAgent(
             config=self.config, dataset_path=self.task_path, workspace=self.workspace
         ).create_graph()
-        researcher_agent = ResearcherAgent(
-            inference=self.inference, embeddings=self.embeddings
-        ).create_graph()
+        researcher_agent = ResearcherAgent(config=self.config).create_graph()
 
         entry_point = SupervisorAgent(
-            inference=self.inference,
+            config=self.config,
             automl_agent=automl_agent,
             researcher_agent=researcher_agent,
         ).create_graph()
@@ -165,9 +164,7 @@ async def ask(self, message: str) -> AsyncIterator[Any]:
             f"FedotAI ask called with message (first 100 chars): '{message[:100]}...'"
         )
         if not self.workspace:
-            self.workspace = Path(
-                f"fedotllm-output-{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}"
-            )
+            self.workspace = Path(f"fedotllm-output-{self.config.session_id}")
             logger.info(f"Workspace created at: {self.workspace}")
         translator_agent = TranslatorAgent(inference=self.inference)
         logger.info("Translating input message to English for ask.")
@@ -179,11 +176,9 @@ async def ask(self, message: str) -> AsyncIterator[Any]:
         automl_agent = AutoMLAgent(
             config=self.config, dataset_path=self.task_path, workspace=self.workspace
         ).create_graph()
-        researcher_agent = ResearcherAgent(
-            inference=self.inference, embeddings=self.embeddings
-        ).create_graph()
+        researcher_agent = ResearcherAgent(config=self.config).create_graph()
         entry_point = SupervisorAgent(
-            inference=self.inference,
+            config=self.config,
             automl_agent=automl_agent,
             researcher_agent=researcher_agent,
         ).create_graph()
diff --git a/fedotllm/prompts/automl.py b/fedotllm/prompts/automl.py
@@ -60,7 +60,11 @@ def fix_solution_prompt(
 {stderr}
 ```
 
-Write the full fixed code below.
+**IMPORTANT**:
+- Apply only **minimal**, **targeted** changes to fix code.
+- Do not rewrite from scratch.
+- Write full fixed code solution from file start to end.
+
 ```python
 """
 
diff --git a/fedotllm/prompts/utils.py b/fedotllm/prompts/utils.py
@@ -21,6 +21,30 @@ def structured_response(response_model: Type[T]):
 """
 
 
+def fix_structured_response(
+    json_obj: dict | None, validation_error: str, response_model: Type[T]
+):
+    return f"""
+Correct JSON object to follow Pydantic definitions:
+=====
+```JSON object
+{json_obj}
+```Pydantic definition
+```
+{response_model.model_json_schema()}
+```
+```Validation error
+{validation_error}
+```
+=====
+Important:
+1. Return only valid JSON. No extra explanations, text, or comments.
+2. Ensure that the output can be parsed by a JSON parser directly.
+3. Do not include any non-JSON text or formatting outside the JSON object.
+4. An example is \{{"<object_field>": "<correct_value_for_the_field>"\}}
+"""
+
+
 def ai_assert_prompt(var1, var2, condition: str):
     return f"""
 You are an intelligent assertion function to evaluate conditions between two variables.
diff --git a/fedotllm/settings/config.toml b/fedotllm/settings/config.toml
diff --git a/fedotllm/settings/config_loader.py b/fedotllm/settings/config_loader.py
diff --git a/fedotllm/web/backend/app.py b/fedotllm/web/backend/app.py
@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Literal, Optional, List
+from typing import List, Literal, Optional
 
 from typing_extensions import AsyncIterator
 
diff --git a/fedotllm/web/frontend/pages/chat.py b/fedotllm/web/frontend/pages/chat.py
diff --git a/fedotllm/web/frontend/pages/init_session.py b/fedotllm/web/frontend/pages/init_session.py
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/tests/unit/agents/test_translator.py b/tests/unit/agents/test_translator.py
diff --git a/tests/unit/llm/test_inference.py b/tests/unit/llm/test_inference.py
