Skip to content

Commit 57237f0

Browse files
authored
feat: ✨ Allow customization of LLM question (#252)
* feat: ✨ Allow customization of LLM question * fix env variable name * Add an option to get the LLM question from a file * Add a prefix in the ENVs to diambiguate them * fix property access
1 parent c663a9b commit 57237f0

File tree

6 files changed

+32
-11
lines changed

6 files changed

+32
-11
lines changed

README.md

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ By default, there is a `GenericProvider` that supports a `SimpleProcessor` using
8888
8989
#### LLM-powered Parsers
9090

91-
The library supports an optional parser option leveraging Large Language Model (LLM) to provide best-effort parsing when the specific parsers have not been successful.
91+
The library supports an optional parser option leveraging Large Language Models (LLM) to provide best-effort parsing when the specific parsers have not been successful.
9292

9393
> Warning: Some of these integrations, such as OpenAI, require of extras installations parameters. Check the [extras section](#extras)
9494
@@ -98,9 +98,12 @@ When the appropriate environment variable(s) are set (see below), these LLM pars
9898
9999
These are the currently supported LLM integrations:
100100

101+
- `PARSER_LLM_QUESTION_STR` (Optional), question to overwrite the default one. Change it carefully. It has precedence over `PARSER_LLM_QUESTION_FILEPATH`
102+
- `PARSER_LLM_QUESTION_FILEPATH` (Optional), a path to a file that contains a question to overwrite the default one.
103+
101104
- [OpenAI](https://openai.com/product), these are the supported ENVs:
102-
- `OPENAI_API_KEY` (Required): OpenAI API Key.
103-
- `OPENAI_MODEL` (Optional): The LLM model to use, defaults to "gpt-3.5-turbo".
105+
- `PARSER_OPENAI_API_KEY` (Required): OpenAI API Key.
106+
- `PARSER_OPENAI_MODEL` (Optional): The LLM model to use, defaults to "gpt-3.5-turbo".
104107

105108
### Metadata
106109

circuit_maintenance_parser/parser.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""Definition of Mainentance Notification base classes."""
22
import logging
3+
import os
34
import base64
45
import calendar
56
import datetime
@@ -346,6 +347,23 @@ def get_key_with_string(dictionary: dict, string: str):
346347
return key
347348
return None
348349

350+
@property
351+
def llm_question(self):
352+
"""Return the LLM question."""
353+
custom_llm_question = os.getenv("PARSER_LLM_QUESTION_STR")
354+
if custom_llm_question:
355+
return custom_llm_question
356+
357+
custom_llm_question_path = os.getenv("PARSER_LLM_QUESTION_FILEPATH")
358+
if custom_llm_question_path:
359+
try:
360+
with open(custom_llm_question_path, mode="r", encoding="utf-8") as llm_question_file:
361+
return llm_question_file.read()
362+
except OSError as err:
363+
logger.warning("The file %s can't be read: %s", custom_llm_question_path, err)
364+
365+
return self._llm_question
366+
349367
def get_llm_response(self, content):
350368
"""Method to retrieve the response from the LLM for some content."""
351369
raise NotImplementedError

circuit_maintenance_parser/parsers/openai.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,15 +24,15 @@ def get_llm_response(self, content) -> Optional[List]:
2424
if not _HAS_OPENAI:
2525
raise ImportError("openai extra is required to use OpenAIParser.")
2626

27-
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
28-
model = os.getenv("OPENAI_MODEL", "gpt-3.5-turbo")
27+
client = OpenAI(api_key=os.getenv("PARSER_OPENAI_API_KEY"))
28+
model = os.getenv("PARSER_OPENAI_MODEL", "gpt-3.5-turbo")
2929
try:
3030
response = client.chat.completions.create(
3131
model=model,
3232
messages=[
3333
{ # type: ignore
3434
"role": "system",
35-
"content": self._llm_question,
35+
"content": self.llm_question,
3636
},
3737
{ # type: ignore
3838
"role": "user",

circuit_maintenance_parser/provider.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ def get_maintenances(self, data: NotificationData) -> Iterable[Maintenance]:
123123
logger.debug("Skipping notification %s due filtering policy for %s.", data, self.__class__.__name__)
124124
return []
125125

126-
if os.getenv("OPENAI_API_KEY"):
126+
if os.getenv("PARSER_OPENAI_API_KEY"):
127127
self._processors.append(CombinedProcessor(data_parsers=[EmailDateParser, OpenAIParser]))
128128

129129
for processor in self._processors:

tests/conftest.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,16 @@
22
import os
33

44

5-
token_openai = os.getenv("OPENAI_API_KEY")
5+
token_openai = os.getenv("PARSER_OPENAI_API_KEY")
66

77

88
def pytest_configure(config): # pylint: disable=unused-argument
99
"""Clean environment for tests."""
1010
if token_openai:
11-
del os.environ["OPENAI_API_KEY"]
11+
del os.environ["PARSER_OPENAI_API_KEY"]
1212

1313

1414
def pytest_sessionfinish(session, exitstatus): # pylint: disable=unused-argument
1515
"""Recove environment after tests."""
1616
if token_openai:
17-
os.environ["OPENAI_API_KEY"] = token_openai
17+
os.environ["PARSER_OPENAI_API_KEY"] = token_openai

tests/unit/test_providers.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ class ProviderWithIncludeFilter(GenericProvider):
117117
)
118118
def test_provider_gets_mlparser(provider_class):
119119
"""Test to check the any provider gets a default ML parser when ENV is activated."""
120-
os.environ["OPENAI_API_KEY"] = "some_api_key"
120+
os.environ["PARSER_OPENAI_API_KEY"] = "some_api_key"
121121
data = NotificationData.init_from_raw("text/plain", b"fake data")
122122
data.add_data_part("text/html", b"other data")
123123

0 commit comments

Comments
 (0)