Добавлен вывод сгенерированных данных

build-dataset/config.yml

@@ -0,0 +1,53 @@
+llm:
+  provider: codex  # openai, codex
+  # model: gpt-6-Astra
+  # model: qwen/qwen3-8b
+
+generation:
+  temperature: 0.7
+  max_tokens: 4096
+
+runtime:
+  concurrency: 1 # Колличество одновременных запросов
+
+prompts:
+  instructions: |
+    Сформируй последовательный и логически корректный план выполнения
+    заданной задачи.
+
+    Соблюдай следующие правила:
+    - один шаг плана должен содержать одно элементарное действие;
+    - действия должны выполняться в физически корректной последовательности;
+    - не объединяй несколько действий в один шаг;
+    - учитывай необходимость безопасного перемещения робота;
+    - перед горизонтальным перемещением с объектом при необходимости
+      сначала подними объект;
+    - перед захватом объекта робот должен переместиться к объекту;
+    - перед освобождением объекта робот должен переместиться в точку
+      его размещения;
+    - не добавляй действия, которые не требуются для выполнения задачи;
+    - не добавляй пояснения, комментарии или рассуждения.
+
+    Используй команды в естественной форме, например:
+    - Переместиться к точке <точка>
+    - Подняться вверх на <расстояние> мм
+    - Опуститься вниз на <расстояние> мм
+    - Захватить <объект>
+    - Освободить <объект>
+    - Повернуть захват на <угол> градусов
+    - Переместиться в исходное положение
+
+    Если необходимые параметры явно указаны в задании, используй их.
+    Не придумывай неизвестные координаты, расстояния, точки или другие
+    параметры, отсутствующие в задании.
+  system: |
+    Ты являешься системой высокоуровневого планирования действий
+    коллаборативного робота.
+
+    Твоя задача — формировать последовательность высокоуровневых
+    команд управления роботом для выполнения заданной технологической
+    операции.
+  user: |
+    Сформируй высокоуровневый план выполнения следующей задачи:
+
+    {task}

build-dataset/robot-move-data.py

@@ -1,5 +1,7 @@
 import asyncio
+from pathlib import Path
 
+import pandas as pd
 from dotenv import load_dotenv
 
 # from pydantic import BaseModel
@@ -8,14 +10,62 @@ from pydantic_ai.settings import ModelSettings
 
 from llm.config import load_config
 from llm.factory import LLMModelProvider
+from llm.schemas import DatasetRecord, Message, RobotPlan, SFTFormat
 
 load_dotenv()
 
 
+def save_to_csv(record: DatasetRecord,
+                path: str | Path):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    exists = path.exists()
+
+    df = pd.DataFrame([record.model_dump()])
+    df.to_csv(path,
+              mode="a" if exists else "w",
+              header=not exists,
+              index=False,
+              encoding="utf-8")
+
+def save_to_jsonl(sft_format: SFTFormat,
+                  path: str | Path):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    
+    with path.open("a", encoding="utf-8") as file:
+        file.write(sft_format.model_dump_json())
+        file.write("\n")
+
+
+def record_to_sft(record: DatasetRecord) -> SFTFormat:
+    return SFTFormat(
+        message=[
+            Message(
+                role="system",
+                content=record.system_prompt
+            ),
+            Message(
+                role="user",
+                content=record.user_prompt
+            ),
+            Message(
+                role="assistant",
+                content=record.assistant_content
+            )
+        ]
+    )
+
+def format_robot_plan(plan: RobotPlan):
+    return "\n".join(
+        f"{index}. {step}"
+        for index, step in enumerate(plan.steps, start=1)
+    )
+
+
 async def main():
-
-
-    config = load_config("config.yml")
+    config = load_config("./build-dataset/config.yml")
     model = await LLMModelProvider.create_from_config(config.llm)
 
 
@@ -24,16 +74,37 @@ async def main():
     #                                                   api_key=os.getenv("API_KEY"),
     #                                                   base_url=os.getenv("BASE_URL"))
     agent = Agent(model=model,
+                  output_type=RobotPlan,
                   system_prompt=config.prompts.system,
                   instructions=config.prompts.instructions)
     settings = ModelSettings(
         **config.generation.model_dump()
     )
 
-    user_prompt = config.prompts.user.format(task="Повтори сообщение")
+    task="Перемести деталь A из точки PICK_A в точку PLACE_A. Безопасная высота подъема — 200 мм."
+    user_prompt = config.prompts.user.format(task=task)
     result = await agent.run(user_prompt,
                              model_settings=settings)
-    print(result.output)
+
+    assistant_content = format_robot_plan(result.output)
+    record = DatasetRecord(
+        task=task,
+        system_prompt=config.prompts.system,
+        instructions=config.prompts.instructions,
+        user_prompt=user_prompt,
+        model_response=result.output.model_dump_json(),
+        assistant_content=assistant_content,
+        provider=config.llm.provider,
+        model=config.llm.model,
+        temperature=config.generation.temperature,
+        max_tokens=config.generation.max_tokens
+    )
+
+    save_to_csv(record=record,
+                path="./dataset/robot-move.csv")
+    sft_format = record_to_sft(record)
+    save_to_jsonl(sft_format=sft_format,
+                  path="./dataset/robot-move.jsonl")
 
 
 if __name__ == "__main__":