151 lines
4.6 KiB
Python
151 lines
4.6 KiB
Python
import os
|
|
import webbrowser
|
|
import pandas as pd
|
|
from pathlib import Path
|
|
|
|
from pydantic_ai.models.openai import OpenAIChatModel
|
|
from pydantic_ai.models.openai_codex import OpenAICodexModel
|
|
from pydantic_ai.providers.openai import OpenAIProvider
|
|
from pydantic_ai.providers.openai_codex import (
|
|
OpenAICodexOAuthFlow,
|
|
OpenAICodexProvider,
|
|
)
|
|
|
|
from .config import CodexConfig, LLMConfig, OpenAiConfig
|
|
from .credentials import FileCodexCredentionalSource
|
|
from .schemas import BaseDatasetRecord, ChatRecord, ChatMessage
|
|
|
|
|
|
class LLMModelProvider:
|
|
|
|
@staticmethod
|
|
def create_openai_compatible(model_name: str,
|
|
api_key: str,
|
|
base_url: str
|
|
) -> OpenAICodexModel:
|
|
|
|
provider = OpenAIProvider(base_url=base_url,
|
|
api_key=api_key)
|
|
|
|
return OpenAIChatModel(model_name=model_name,
|
|
provider=provider)
|
|
|
|
@staticmethod
|
|
async def create_codex(model_name: str='gpt-5.6-luna',
|
|
credentional_path: str | Path | None=None
|
|
) -> OpenAICodexModel:
|
|
|
|
if credentional_path is None:
|
|
credentional_path = Path.cwd() / ".config" / "codex-credentials.json"
|
|
|
|
|
|
source = FileCodexCredentionalSource(credentional_path)
|
|
|
|
if not source.path.exists():
|
|
flow = OpenAICodexOAuthFlow()
|
|
webbrowser.open(flow.authorization_url())
|
|
credentials = await flow.exchange_code_from_callback()
|
|
await source.save(credentials)
|
|
|
|
provider = OpenAICodexProvider(credential_source=source)
|
|
|
|
return OpenAICodexModel(model_name=model_name,
|
|
provider=provider)
|
|
|
|
@staticmethod
|
|
async def create_from_config(config: LLMConfig):
|
|
|
|
if isinstance(config, CodexConfig):
|
|
return await LLMModelProvider.create_codex(
|
|
model_name=config.model,
|
|
credentional_path=config.credential_path
|
|
)
|
|
|
|
if isinstance(config, OpenAiConfig):
|
|
api_key = os.getenv(config.api_key_env)
|
|
base_url=os.getenv(config.base_url_env)
|
|
|
|
if not api_key:
|
|
raise ValueError(
|
|
f"Environment variable "
|
|
f"{config.api_key_env!r} is not set"
|
|
)
|
|
|
|
if not base_url:
|
|
raise ValueError(
|
|
f"Environment variable "
|
|
f"{config.base_url_env!r} is not set"
|
|
)
|
|
|
|
return LLMModelProvider.create_openai_compatible(
|
|
model_name=config.model,
|
|
api_key=api_key,
|
|
base_url=base_url
|
|
)
|
|
|
|
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") |