Merge pull request #2 from shohey1226/schema_with_funcation_calling

Schema with funcation calling

lib/llm_memory/broca.rb

@@ -43,6 +43,46 @@ def respond(args)
       end
     end
 
+    def respond_with_schema(context: {}, schema: {})
+      response_content = respond(context)
+      begin
+        response = client.chat(
+          parameters: {
+            model: "gpt-3.5-turbo-0613", # as of July 3, 2023
+            messages: [
+              {
+                role: "user",
+                content: response_content
+              }
+            ],
+            functions: [
+              {
+                name: "broca",
+                description: "Formating the content with the specified schema",
+                parameters: schema
+              }
+            ]
+          }
+        )
+        LlmMemory.logger.debug(response)
+        message = response.dig("choices", 0, "message")
+        if message["role"] == "assistant" && message["function_call"]
+          function_name = message.dig("function_call", "name")
+          args =
+            JSON.parse(
+              message.dig("function_call", "arguments"),
+              {symbolize_names: true}
+            )
+          if function_name == "broca"
+            args
+          end
+        end
+      rescue => e
+        LlmMemory.logger.info(e.inspect)
+        nil
+      end
+    end
+
     def generate_prompt(args)
       erb = ERB.new(@prompt)
       erb.result_with_hash(args)

spec/llm_memory/broca_spec.rb

@@ -47,5 +47,24 @@
       res = broca.respond(related_docs: related_docs, query_str: "what is my name?")
       expect(res).to include("Shohei")
     end
+
+    it "runs respond_with_schema method", :vcr do
+      related_docs = [{content: "My name is Shohei"}, {content: "I'm a software engineer"}]
+      broca = LlmMemory::Broca.new(prompt: template)
+      res = broca.respond_with_schema(
+        context: {related_docs: related_docs, query_str: "what is my name?"},
+        schema: {
+          type: :object,
+          properties: {
+            name: {
+              type: :string,
+              description: "The name of person"
+            }
+          },
+          required: ["name"]
+        }
+      )
+      expect(res).to include({name: "Shohei"})
+    end
   end
 end