[Python] Pyre Pysa LLM calls as source and sink

I'm trying to define LLM calls as both sinks and sources inside Pysa (from Pyre). I managed to define other sinks and sources (like input, os.system, etc), but can't figure that one out. I'm testing on the following program. Ideally, I don't want to make any modifications to the analyzed repos and treat them as "read only":

import random
import openai

client = openai.Client()

def func1(x):
func2(x)

def func2(x):
client.responses.create(
model="gpt-4.1",
input=x
)

def vulnerable():
x = client.responses.create(
model="gpt-4.1",
input="Write a one-sentence bedtime story about a unicorn."
)

if random.random() > 0.5:
func2(x)
else:
func1(x)

if __name__ == "__main__":
vulnerable()


One example of a failing pysa file and taint.config (I tried a lot) ...

llms.pysa:

# anything.chat.completions.create(...)
def __.chat.completions.create(
self,
*,
model,
messages: TaintSink[LLMInput],
**kwargs
) -> TaintSource[LLMOutput]: ...

# anything.responses.create(...)
def __.responses.create(
self,
*,
model,
input: TaintSink[LLMInput],
**kwargs
) -> TaintSource[LLMOutput]: ...

# anything.messages.create(...)
def __.messages.create(
self,
*,
model,
messages: TaintSink[LLMInput],
**kwargs
) -> TaintSource[LLMOutput]: ...


taint.config:

{
"comment": "Pysa taint config with LLM sources/sinks",
"sources": [
{ "name": "UserControlled", "comment": "builtins.input, etc." },
{ "name": "LLMOutput", "comment": "Values returned by any chat.responses.create / chat.completions.create / messages.create" }
],
"sinks": [
{ "name": "RemoteCodeExecution", "comment": "os.system, eval, etc." },
{ "name": "LLMInput", "comment": "Data sent into any chat.responses.create / chat.completions.create / messages.create" }
],
"features": [],
"rules": [
{
"name": "Shell injection",
"code": 5001,
"sources": ["UserControlled"],
"sinks": ["RemoteCodeExecution"],
"message_format": "User input might reach os.system"
},
{
"name": "LLM invocation",
"code": 5002,
"sources": ["UserControlled"],
"sinks": ["LLMInput"],
"message_format": "User input might reach an LLM API call"
},
{
"name": "LLM-to-shell",
"code": 5003,
"sources": ["LLMOutput"],
"sinks": ["RemoteCodeExecution"],
"message_format": "LLM output might reach os.system"
}
]
}


The error I get when running pyre validate-models:

ƛ Found 3 model verification errors!
taint/llms.pysa:2:0 `__.chat.completions.create` is not part of the environment, no module `__` in search path.
taint/llms.pysa:11:0 `__.responses.create` is not part of the environment, no module `__` in search path.
taint/llms.pysa:20:0 `__.messages.create` is not part of the environment, no module `__` in search path.


Ideally, I want to support calls to OpenAI and Anthropic, which are (they may have several more kwargs):

response = client.chat.completions.create(
model="gpt-4.1",
messages=[{"role": role_name, "content": "Hello OpenAI."}, {"role": "user", "content":
"Bye OpenAI."}]
)

response = client.responses.create(
model="gpt-4.1",
input="Hello OpenAI."
)

response = client.messages.create(
model="claude-opus-4-20250514",
max_tokens=1024,
messages=[ {"role": "user", "content": "Hello, Claude"} ]
)


I would appreciate help a lot, thanks!

Continue reading...