Skip to content

Commit 47b0eb1

Browse files
authored
Fix temp file resource leak (#569)
ChatOpenAI will soon be consolidated, but this fixes a resource leak within it. Also adds more helpful type hinting on File.
1 parent b9e32cc commit 47b0eb1

2 files changed

Lines changed: 36 additions & 31 deletions

File tree

src/steamship/agents/llms/openai.py

Lines changed: 35 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -91,35 +91,40 @@ def chat(self, messages: List[Block], tools: Optional[List[Tool]], **kwargs) ->
9191
tags=[Tag(kind=TagKind.GENERATION, name=GenerationTag.PROMPT_COMPLETION)],
9292
)
9393

94-
options = {}
95-
if len(tools) > 0:
96-
functions = []
97-
for tool in tools:
98-
functions.append(tool.as_openai_function().dict())
99-
options["functions"] = functions
100-
101-
if "max_tokens" in kwargs:
102-
options["max_tokens"] = kwargs["max_tokens"]
103-
104-
extra = {
105-
AgentLogging.LLM_NAME: "OpenAI",
106-
AgentLogging.IS_MESSAGE: True,
107-
AgentLogging.MESSAGE_TYPE: AgentLogging.PROMPT,
108-
AgentLogging.MESSAGE_AUTHOR: AgentLogging.LLM,
109-
}
110-
111-
if logging.DEBUG >= logging.root.getEffectiveLevel():
112-
extra["messages"] = json.dumps(
113-
"\n".join([f"[{msg.chat_role}] {msg.as_llm_input()}" for msg in messages])
94+
try:
95+
options = {}
96+
if len(tools) > 0:
97+
functions = []
98+
for tool in tools:
99+
functions.append(tool.as_openai_function().dict())
100+
options["functions"] = functions
101+
102+
if "max_tokens" in kwargs:
103+
options["max_tokens"] = kwargs["max_tokens"]
104+
105+
extra = {
106+
AgentLogging.LLM_NAME: "OpenAI",
107+
AgentLogging.IS_MESSAGE: True,
108+
AgentLogging.MESSAGE_TYPE: AgentLogging.PROMPT,
109+
AgentLogging.MESSAGE_AUTHOR: AgentLogging.LLM,
110+
}
111+
112+
if logging.DEBUG >= logging.root.getEffectiveLevel():
113+
extra["messages"] = json.dumps(
114+
"\n".join([f"[{msg.chat_role}] {msg.as_llm_input()}" for msg in messages])
115+
)
116+
extra["tools"] = ",".join([t.name for t in tools])
117+
else:
118+
extra["num_messages"] = len(messages)
119+
extra["num_tools"] = len(tools)
120+
121+
logging.info(f"OpenAI ChatComplete ({messages[-1].as_llm_input()})", extra=extra)
122+
123+
tool_selection_task = self.generator.generate(
124+
input_file_id=temp_file.id, options=options
114125
)
115-
extra["tools"] = ",".join([t.name for t in tools])
116-
else:
117-
extra["num_messages"] = len(messages)
118-
extra["num_tools"] = len(tools)
119-
120-
logging.info(f"OpenAI ChatComplete ({messages[-1].as_llm_input()})", extra=extra)
121-
122-
tool_selection_task = self.generator.generate(input_file_id=temp_file.id, options=options)
123-
tool_selection_task.wait()
126+
tool_selection_task.wait()
124127

125-
return tool_selection_task.output.blocks
128+
return tool_selection_task.output.blocks
129+
finally:
130+
temp_file.delete()

src/steamship/data/file.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ def create(
121121
blocks: List[Block] = None,
122122
tags: List[Tag] = None,
123123
public_data: bool = False,
124-
) -> Any:
124+
) -> File:
125125

126126
req = {
127127
"handle": handle,

0 commit comments

Comments
 (0)