chore: adjust default parameters and test assertion

This commit is contained in:
100gle
2023-09-27 17:01:14 +08:00
parent 8477774b62
commit 77201d5de2
2 changed files with 12 additions and 5 deletions

View File

@@ -44,16 +44,16 @@ class OpenAIParser:
def __init__(
self,
api_key: str,
api_base: str | None = None,
model: str | None = None,
api_base: str = "https://api.openai.com/v1",
model: str = "gpt-3.5-turbo",
**kwargs,
) -> None:
if not api_key:
raise ValueError("API key is required.")
self._api_key = api_key
self.api_base = api_base or "https://api.openai.com/v1"
self.model = model or "gpt-3.5-turbo"
self.api_base = api_base
self.model = model
self.openai_kwargs = kwargs
def parse(self, text: str, prompt: str | None = None) -> str:

View File

@@ -1,3 +1,4 @@
import json
import os
from dotenv import load_dotenv
@@ -15,4 +16,10 @@ class TestOpenAIParser:
def test_parse(self):
text = "[梦蓝字幕组]New Doraemon 哆啦A梦新番[747][2023.02.25][AVC][1080P][GB_JP][MP4]"
result = self.parser.parse(text=text)
assert result == "XXX"
assert json.loads(result) == {
"group": "梦蓝字幕组",
"title_en": "New Doraemon",
"resolution": "1080P",
"episode": 747,
"season": 1,
}