From a382e8f721f3407747fae26c3d1e301e45a912bd Mon Sep 17 00:00:00 2001
From: Oded Shopen <odedia@gmail.com>
Date: Tue, 10 Sep 2024 21:44:44 +0300
Subject: [PATCH] support for function calling on vets

---
 .../petclinic/genai/AIFunctionConfiguration.java | 16 ++++++++++++++++
 .../petclinic/genai/PetclinicAiProvider.java     | 13 ++++++++++++-
 .../application-postgres-openai.properties       |  2 +-
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/springframework/samples/petclinic/genai/AIFunctionConfiguration.java b/src/main/java/org/springframework/samples/petclinic/genai/AIFunctionConfiguration.java
index ea331c0ed..f46fd8ea3 100644
--- a/src/main/java/org/springframework/samples/petclinic/genai/AIFunctionConfiguration.java
+++ b/src/main/java/org/springframework/samples/petclinic/genai/AIFunctionConfiguration.java
@@ -8,6 +8,7 @@ import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Description;
 import org.springframework.samples.petclinic.owner.Owner;
 import org.springframework.samples.petclinic.owner.Pet;
+import org.springframework.samples.petclinic.vet.Vet;
 
 @Configuration
 class AIFunctionConfiguration {
@@ -20,6 +21,14 @@ class AIFunctionConfiguration {
 		};
 	}
 
+	@Bean
+	@Description("List the veterinarians that the pet clinic has")
+	public Function<VetRequest, VetResponse> listVets(PetclinicAiProvider petclinicAiProvider) {
+		return request -> {
+			return petclinicAiProvider.getAllVets();
+		};
+	}
+
 	@Bean
 	@Description("Add a pet to an owner identified by the ownerId")
 	public Function<AddPetRequest, AddedPetResponse> addPetToOwner(PetclinicAiProvider petclinicAiProvider) {
@@ -41,3 +50,10 @@ record OwnersResponse(List<Owner> owners) {
 
 record AddedPetResponse(Owner owner) {
 };
+
+record VetResponse(List<Vet> vet) {
+};
+
+record VetRequest(Vet vet) {
+
+}
\ No newline at end of file
diff --git a/src/main/java/org/springframework/samples/petclinic/genai/PetclinicAiProvider.java b/src/main/java/org/springframework/samples/petclinic/genai/PetclinicAiProvider.java
index 6ba98d0df..364c20659 100644
--- a/src/main/java/org/springframework/samples/petclinic/genai/PetclinicAiProvider.java
+++ b/src/main/java/org/springframework/samples/petclinic/genai/PetclinicAiProvider.java
@@ -5,6 +5,8 @@ import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Pageable;
 import org.springframework.samples.petclinic.owner.Owner;
 import org.springframework.samples.petclinic.owner.OwnerRepository;
+import org.springframework.samples.petclinic.vet.Vet;
+import org.springframework.samples.petclinic.vet.VetRepository;
 import org.springframework.stereotype.Service;
 
 @Service
@@ -12,8 +14,11 @@ public class PetclinicAiProvider {
 
 	OwnerRepository ownerRepository;
 
-	public PetclinicAiProvider(OwnerRepository ownerRepository) {
+	VetRepository vetRepository;
+
+	public PetclinicAiProvider(OwnerRepository ownerRepository, VetRepository vetRepository) {
 		this.ownerRepository = ownerRepository;
+		this.vetRepository = vetRepository;
 	}
 
 	public OwnersResponse getAllOwners() {
@@ -22,6 +27,12 @@ public class PetclinicAiProvider {
 		return new OwnersResponse(ownerPage.getContent());
 	}
 
+	public VetResponse getAllVets() {
+		Pageable pageable = PageRequest.of(0, Integer.MAX_VALUE);
+		Page<Vet> vetsPage = vetRepository.findAll(pageable);
+		return new VetResponse(vetsPage.getContent());
+	}
+
 	public AddedPetResponse addPetToOwner(AddPetRequest request) {
 		Owner owner = ownerRepository.findById(request.ownerId());
 		owner.addPet(request.pet());
diff --git a/src/main/resources/application-postgres-openai.properties b/src/main/resources/application-postgres-openai.properties
index 6ec85c0f0..587daad8b 100644
--- a/src/main/resources/application-postgres-openai.properties
+++ b/src/main/resources/application-postgres-openai.properties
@@ -5,5 +5,5 @@ spring.datasource.username=${POSTGRES_USER:petclinic}
 spring.datasource.password=${POSTGRES_PASS:petclinic}
 # SQL is written to be idempotent so this is safe
 spring.sql.init.mode=always
-spring.ai.azure.openai.chat.options.functions=listOwners,addPetToOwner
+spring.ai.azure.openai.chat.options.functions=listOwners,listVets,addPetToOwner
 spring.ai.azure.openai.chat.options.temperature: 0.7