diff --git a/trustgraph-base/trustgraph/clients/llm_client.py b/trustgraph-base/trustgraph/clients/llm_client.py index 3b52cb16..cfb0e606 100644 --- a/trustgraph-base/trustgraph/clients/llm_client.py +++ b/trustgraph-base/trustgraph/clients/llm_client.py @@ -35,6 +35,8 @@ class LlmClient(BaseClient): output_schema=TextCompletionResponse, ) - def request(self, prompt, timeout=300): - return self.call(prompt=prompt, timeout=timeout).response + def request(self, system, prompt, timeout=300): + return self.call( + system=system, prompt=prompt, timeout=timeout + ).response diff --git a/trustgraph-base/trustgraph/schema/models.py b/trustgraph-base/trustgraph/schema/models.py index 2196a3d2..70cb2c8f 100644 --- a/trustgraph-base/trustgraph/schema/models.py +++ b/trustgraph-base/trustgraph/schema/models.py @@ -9,6 +9,7 @@ from . types import Error # LLM text completion class TextCompletionRequest(Record): + system = String() prompt = String() class TextCompletionResponse(Record): diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index e2f31c52..a9c05cc8 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -101,7 +101,7 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) - prompt = v.prompt + prompt = v.system + "\n\n" + v.prompt try: diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py index eea6b8c4..2e5416f4 100755 --- a/trustgraph-flow/trustgraph/model/prompt/template/service.py +++ b/trustgraph-flow/trustgraph/model/prompt/template/service.py @@ -146,7 +146,7 @@ class Processor(ConsumerProducer): def request(self, system, prompt): print(system) print(prompt, flush=True) - return self.llm.request(system + "\n\n" + prompt) + return self.llm.request(system, prompt) self.llm = Llm(self.llm) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py index cbc638ec..4db7dbf1 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py @@ -136,7 +136,7 @@ class Processor(ConsumerProducer): try: prompt = self.build_prompt( - "You are a helpful chatbot", + v.system, v.prompt ) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py index 1dbfba27..a3edb859 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -95,7 +95,7 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) - prompt = v.prompt + prompt = v.system + "\n\n" + v.prompt try: diff --git a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py index d784f1f8..01ce837d 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -95,7 +95,7 @@ class Processor(ConsumerProducer): model=self.model, max_tokens=self.max_output, temperature=self.temperature, - system = "You are a helpful chatbot.", + system = v.system, messages=[ { "role": "user", diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index 8b72f9db..d03e1554 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -79,6 +79,7 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) + system = v.system prompt = v.prompt try: @@ -88,7 +89,7 @@ class Processor(ConsumerProducer): output = self.cohere.chat( model=self.model, message=prompt, - preamble = "You are a helpful AI-assistant.", + preamble = system, temperature=self.temperature, chat_history=[], prompt_truncation='auto', diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index 01dc9736..a249998d 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -111,7 +111,14 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) - prompt = v.prompt + # FIXME: There's a system prompt above. Maybe if system changes, + # then reset self.llm? It shouldn't do, because system prompt + # is set system wide? + + # Or... could keep different LLM structures for different system + # prompts? + + prompt = v.system + "\n\n" + v.prompt try: diff --git a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py index a25ad8ec..274948a8 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py @@ -84,7 +84,7 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) - prompt = v.prompt + prompt = v.system + "\n\n" + v.prompt try: diff --git a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py index 17151a00..00d44f6d 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py @@ -80,7 +80,7 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) - prompt = v.prompt + prompt = v.system + "\n\n" + v.prompt try: diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index 849e78ea..b1fe7c94 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py @@ -83,7 +83,7 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) - prompt = v.prompt + prompt = v.system + "\n\n" + v.prompt try: diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index ecdbae06..cb817836 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -143,7 +143,7 @@ class Processor(ConsumerProducer): print(f"Handling prompt {id}...", flush=True) - prompt = v.prompt + prompt = v.system + "\n\n" + v.prompt with __class__.text_completion_metric.time():