Implement system in text completion API (#137)

* Add system prompt to LLM invocation

* Added system parameter to LLMs

* Added to Bedrock and VertexAI

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
 

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):

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:
 

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)
 

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
             )
 

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:

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",

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',

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:
 

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:
 

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:
 

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:
 

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():