greek_lang/src/greek_lang/openai_manager/manager.py
2025-06-22 14:37:27 +03:00

88 lines
2.7 KiB
Python

from __future__ import annotations
import dataclasses
import pydantic
from openai import AsyncOpenAI
from greek_lang.languages import LanguageEnum
from greek_lang.glossaries.models import LexicalCategoryEnum
class WordInfo(pydantic.BaseModel):
lemma: str = pydantic.Field(
...,
description="lemma (base form) - for verbs, use the 1st person singular in present indicative, "
"for nouns and adjectives, use the nominative singular masculine (for adjectives)",
)
transcription: str = pydantic.Field(
...,
description="lemma phonetic transcription in IPA",
)
translation: str = pydantic.Field(
...,
description="lemma translation in {target_language}",
)
description: str = pydantic.Field(
...,
description="lemma description in {target_language}",
)
part_of_speech: str = pydantic.Field(
...,
description=f"part of speech, one of {[cat.value for cat in LexicalCategoryEnum]}",
)
example: str = pydantic.Field(
...,
description="lemma example",
)
example_transcription: str = pydantic.Field(
...,
description="lemma phonetic transcription in IPA of an example",
)
example_translation: str = pydantic.Field(
...,
description="lemma translation of the example in {target_language}",
)
category: str = pydantic.Field(
...,
description=f"semantic category in {{target_language}}",
)
etymology: str = pydantic.Field(
...,
description="short etymology of the word described in {target_language}",
)
@dataclasses.dataclass(frozen=True)
class OpenAiManager:
client: AsyncOpenAI
async def get_gpt_response(
self,
*,
word: str,
source_lang: LanguageEnum,
target_lang: LanguageEnum,
model: str = "gpt-4o",
) -> WordInfo:
system_message = {
"role": "system",
"content": "You are a helpful assistant that provides detailed word information.",
}
user_message = {
"role": "user",
"content": f'Provide detailed information about the word "{word}" in language {source_lang!s}, set {{target_language}} = {target_lang!s}.',
}
response = await self.client.beta.chat.completions.parse(
model=model,
messages=( # type: ignore
system_message,
user_message,
),
response_format=WordInfo,
)
word_info: WordInfo | None = response.choices[0].message.parsed
if word_info is None:
raise RuntimeError("No word_info")
return word_info