Обновил структуру проекта

build-dataset/config.yml

@@ -1,6 +1,6 @@
 llm:
   provider: codex  # openai, codex
-  # model: gpt-6-Astra
+  model: gpt-6-astra
   # model: qwen/qwen3-8b
 
 generation:
@@ -8,7 +8,7 @@ generation:
   max_tokens: 4096
 
 runtime:
-  concurrency: 1 # Колличество одновременных запросов
+  concurrency: 5 # Колличество одновременных запросов
 
 prompts:
   instructions: |

build-dataset/robot-move-data.py

@@ -1,7 +1,5 @@
 import asyncio
-from pathlib import Path
 
-import pandas as pd
 from dotenv import load_dotenv
 
 # from pydantic import BaseModel
@@ -9,58 +7,53 @@ from pydantic_ai import Agent
 from pydantic_ai.settings import ModelSettings
 
 from llm.config import load_config
-from llm.factory import LLMModelProvider
+from llm.factory import LLMModelProvider, DatasetWriter
-from llm.schemas import DatasetRecord, Message, RobotPlan, SFTFormat
+from llm.schemas import (RobotPlan, RobotDatasetRecord, 
+                         Prompts, ModelResponse, 
+                         ModelInfo, GenerationSettings)
 
 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}"
+        f"{step}"
-        for index, step in enumerate(plan.steps, start=1)
+        for step in plan.steps
+    )
+
+
+async def generate_record(task: str,
+                           agent: Agent,
+                           settings: ModelSettings,
+                           config,
+                           semaphore: asyncio.Semaphore) -> RobotDatasetRecord:
+
+    async with semaphore:
+        user_prompt = config.prompts.user.format(task=task)
+    
+    result = await agent.run(user_prompt,
+                             model_settings=settings)
+    assistant_content = format_robot_plan(result.output)
+
+    return RobotDatasetRecord(
+        task=task,
+        prompts=Prompts(
+            system_prompt=config.prompts.system,
+            instructions=config.prompts.instructions,
+            user_prompt=user_prompt
+        ),
+        response=ModelResponse(
+            model_response=result.output.model_dump_json(),
+            assistant_content=assistant_content
+        ),
+        model_info=ModelInfo(
+            provider=config.llm.provider,
+            model=config.llm.model
+        ),
+        generation=GenerationSettings(
+            temperature=config.generation.temperature,
+            max_tokens=config.generation.max_tokens
+        )
     )
     
 
@@ -81,30 +74,42 @@ async def main():
         **config.generation.model_dump()
     )
 
-    task="Перемести деталь A из точки PICK_A в точку PLACE_A. Безопасная высота подъема — 200 мм."
+    tasks = [
-    user_prompt = config.prompts.user.format(task=task)
+        (
-    result = await agent.run(user_prompt,
+            "Перемести деталь A из точки PICK_A в точку PLACE_A. "
-                             model_settings=settings)
+            "Безопасная высота подъема — 200 мм."
+        ),
+        (
+            "Перемести деталь B из точки PICK_B в точку PLACE_B. "
+            "Безопасная высота подъема — 150 мм."
+        ),
+        (
+            "Возьми крышку из точки COVER_PICK и перемести "
+            "ее в точку COVER_PLACE."
+        ),
+    ]
 
-    assistant_content = format_robot_plan(result.output)
+    semaphore = asyncio.Semaphore(config.runtime.concurrency)
-    record = DatasetRecord(
+
-        task=task,
+    records = await asyncio.gather(*[
-        system_prompt=config.prompts.system,
+        generate_record(
-        instructions=config.prompts.instructions,
+            task,
-        user_prompt=user_prompt,
+            agent,
-        model_response=result.output.model_dump_json(),
+            settings,
-        assistant_content=assistant_content,
+            config,
-        provider=config.llm.provider,
+            semaphore,
-        model=config.llm.model,
+        )
-        temperature=config.generation.temperature,
+        for task in tasks
-        max_tokens=config.generation.max_tokens
+    ])
+
+    for record in records:
+        dataset = DatasetWriter(
+            dataset_record=record,
+            base_path="./dataset"
         )
 
-    save_to_csv(record=record,
+        dataset.to_csv(filename="robot-plan.csv")
-                path="./dataset/robot-move.csv")
+        dataset.to_jsonl(filename="robotplan.jsonl")
-    sft_format = record_to_sft(record)
-    save_to_jsonl(sft_format=sft_format,
-                  path="./dataset/robot-move.jsonl")
 
 
 if __name__ == "__main__":

llm/credentials.py

@@ -10,12 +10,12 @@ from pydantic_ai.providers.openai_codex import (
 
 class FileCodexCredentionalSource(OpenAICodexCredentialSource):
     def __init__(self, 
-                 path: Path,
+                 path: str| Path,
                  encoding: str="utf-8",
                  ensure_ascii: bool=False,
                  indent: int=2):
         super().__init__()
-        self._path = path
+        self._path = Path(path)
         self._encoding=encoding
         self._ensure_ascii=ensure_ascii
         self._indent=indent
@@ -32,6 +32,8 @@ class FileCodexCredentionalSource(OpenAICodexCredentialSource):
             )
 
     async def save(self, credentials: OpenAICodexCredentials) -> None:
+        self.path.parent.mkdir(parents=True , exist_ok=True)
+
         self._path.write_text(
             json.dumps(
                 asdict(credentials),

llm/factory.py

@@ -1,5 +1,6 @@
 import os
 import webbrowser
+import pandas as pd
 from pathlib import Path
 
 from pydantic_ai.models.openai import OpenAIChatModel
@@ -12,6 +13,7 @@ from pydantic_ai.providers.openai_codex import (
 
 from .config import CodexConfig, LLMConfig, OpenAiConfig
 from .credentials import FileCodexCredentionalSource
+from .schemas import BaseDatasetRecord, ChatRecord, ChatMessage
 
 
 class LLMModelProvider:
@@ -82,3 +84,68 @@ class LLMModelProvider:
             )
 
         raise ValueError(f"Unsupported LLM config: {type(config)}")
+
+class DatasetWriter:
+    def __init__(self,
+                 dataset_record: BaseDatasetRecord,
+                 base_path: str | Path):
+
+
+        self._dataset_record = dataset_record    
+        self._path = Path(base_path)
+        self._path.mkdir(parents=True, exist_ok=True)
+
+    def __check_file_exists(self, filename :str)-> bool:
+        file_path = self._path / filename
+        return file_path.exists()
+
+    def to_csv(self,
+               filename: str) -> None:
+
+        exists=self.__check_file_exists(filename)
+        
+        df = pd.json_normalize(
+            self._dataset_record.model_dump()
+        )
+
+        df.to_csv(self._path / filename,
+                  mode="a" if exists else "w",
+                  header=not exists,
+                  index=False,
+                  encoding="utf-8")
+
+    def to_chat_record(self) -> ChatRecord:
+        response = self._dataset_record.response
+        prompts = self._dataset_record.prompts
+
+        assistant_content = (
+            response.assistant_content
+            or response.model_response
+        )
+
+        return ChatRecord(
+            messages=[
+                ChatMessage(
+                    role="system",
+                    content=prompts.system_prompt,
+                ),
+                ChatMessage(
+                    role="user",
+                    content=prompts.user_prompt,
+                ),
+                ChatMessage(
+                    role="assistant",
+                    content=assistant_content,
+                ),
+            ]
+        )
+
+    def to_jsonl(self,
+                 filename: str) -> None:
+        
+        chat_record = self.to_chat_record()
+        path = self._path/filename
+
+        with path.open("a", encoding='utf-8') as file:
+            file.write(chat_record.model_dump_json())
+            file.write("\n")

llm/schemas.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Literal, Optional
 
 from pydantic import BaseModel, Field
 
@@ -10,7 +10,7 @@ class RobotPlan(BaseModel):
     )
 
 
-class Message(BaseModel):
+class ChatMessage(BaseModel):
     role: Literal["system",
                   "user",
                   "assistant"]
@@ -18,22 +18,32 @@ class Message(BaseModel):
     content: str
 
 
-class SFTFormat(BaseModel):
+class ChatRecord(BaseModel):
-    message: list[Message]
+    messages: list[ChatMessage]
-
-
-class DatasetRecord(BaseModel):
-    task: str
 
+class Prompts(BaseModel):
     system_prompt: str
     instructions: str
     user_prompt: str
 
-    model_response: str
-    assistant_content: str
 
+class ModelResponse(BaseModel):
+    model_response: str
+    assistant_content: Optional[str] = None
+
+class ModelInfo(BaseModel):
     provider: str
     model: str
 
+class GenerationSettings(BaseModel):
     temperature: float
     max_tokens: int
+
+class BaseDatasetRecord(BaseModel):
+    prompts: Prompts
+    response: ModelResponse
+    model_info: ModelInfo
+    generation: GenerationSettings
+
+class RobotDatasetRecord(BaseDatasetRecord):
+    task: str