diff --git a/devai-api/app/routes.py b/devai-api/app/routes.py index a7e4379..1f974a6 100644 --- a/devai-api/app/routes.py +++ b/devai-api/app/routes.py @@ -28,15 +28,17 @@ from langchain_community.agent_toolkits.gitlab.toolkit import GitLabToolkit from langchain_community.utilities.gitlab import GitLabAPIWrapper from langchain_google_vertexai import ChatVertexAI - +from google.cloud.aiplatform import telemetry from vertexai.generative_models import GenerativeModel +USER_AGENT = 'cloud-solutions/genai-for-developers-v1' model_name="gemini-1.5-pro-preview-0409" -llm = ChatVertexAI(model_name=model_name, - convert_system_message_to_human=True, - temperature=0.2, - max_output_tokens=4096) +with telemetry.tool_context_manager(USER_AGENT): + llm = ChatVertexAI(model_name=model_name, + convert_system_message_to_human=True, + temperature=0.2, + max_output_tokens=4096) gitlab = GitLabAPIWrapper() toolkit = GitLabToolkit.from_gitlab_api_wrapper(gitlab) @@ -63,7 +65,8 @@ async def root(): @routes.get("/test") async def test(): """Test endpoint""" - code_chat = code_chat_model.start_chat() + with telemetry.tool_context_manager(USER_AGENT): + code_chat = code_chat_model.start_chat() response = code_chat.send_message("Tell me about Google Gemini 1.5 capabilities") print(f"Response from Model:\n{response.text}\n") return {"message": response.text} @@ -81,11 +84,11 @@ async def generate_handler(request: Request, prompt: str = Body(embed=True)): REQUIREMENTS: {prompt} """ - - code_chat = code_chat_model.start_chat() + with telemetry.tool_context_manager(USER_AGENT): + code_chat = code_chat_model.start_chat() response = code_chat.send_message(instructions) print(f"Response from Model:\n{response.text}\n") - # GitLab Agent Call + # resp_text = response.candidates[0].content.parts[0].text # pr_prompt = f"""Create GitLab merge request using provided details below. diff --git a/devai-api/requirements.txt b/devai-api/requirements.txt index 94630e7..691e674 100644 --- a/devai-api/requirements.txt +++ b/devai-api/requirements.txt @@ -1,5 +1,5 @@ fastapi==0.109.2 -google-cloud-aiplatform==1.47.0 +google-cloud-aiplatform==1.48.0 google-auth==2.27.0 itsdangerous==2.1.2 jinja2==3.1.3 diff --git a/outer-loop-cli/src/devai/commands/gitlab.py b/outer-loop-cli/src/devai/commands/gitlab.py index 3fae717..6345dfe 100644 --- a/outer-loop-cli/src/devai/commands/gitlab.py +++ b/outer-loop-cli/src/devai/commands/gitlab.py @@ -20,10 +20,11 @@ from langchain_community.utilities.gitlab import GitLabAPIWrapper from google.cloud.aiplatform import telemetry -USER_AGENT = 'cloud-solutions/genai-outer-loop-devai' +USER_AGENT = 'cloud-solutions/genai-for-developers-v1' +model_name="gemini-1.5-pro-preview-0409" with telemetry.tool_context_manager(USER_AGENT): - llm = ChatVertexAI(model_name="gemini-pro", + llm = ChatVertexAI(model_name=model_name, convert_system_message_to_human=True, project=os.environ["PROJECT_ID"], location=os.environ["LOCATION"]) diff --git a/outer-loop-cli/src/devai/commands/jira.py b/outer-loop-cli/src/devai/commands/jira.py index 0c1b8d2..ceec416 100644 --- a/outer-loop-cli/src/devai/commands/jira.py +++ b/outer-loop-cli/src/devai/commands/jira.py @@ -22,12 +22,16 @@ from langchain_community.utilities.jira import JiraAPIWrapper from langchain_google_vertexai import ChatVertexAI from jira import JIRA +from google.cloud.aiplatform import telemetry +USER_AGENT = 'cloud-solutions/genai-for-developers-v1' +model_name="gemini-1.5-pro-preview-0409" -llm = ChatVertexAI(model_name="gemini-pro", - convert_system_message_to_human=True, - project=os.environ["PROJECT_ID"], - location=os.environ["LOCATION"]) +with telemetry.tool_context_manager(USER_AGENT): + llm = ChatVertexAI(model_name=model_name, + convert_system_message_to_human=True, + project=os.environ["PROJECT_ID"], + location=os.environ["LOCATION"]) jira = JiraAPIWrapper() toolkit = JiraToolkit.from_jira_api_wrapper(jira) diff --git a/outer-loop-cli/src/devai/commands/review.py b/outer-loop-cli/src/devai/commands/review.py index d9f1efe..6fd7aee 100644 --- a/outer-loop-cli/src/devai/commands/review.py +++ b/outer-loop-cli/src/devai/commands/review.py @@ -16,17 +16,16 @@ import click from devai.util.file_processor import format_files_as_string from vertexai.generative_models import GenerativeModel, ChatSession +from google.cloud.aiplatform import telemetry + +USER_AGENT = 'cloud-solutions/genai-for-developers-v1' +model_name="gemini-1.5-pro-preview-0409" # Uncomment after configuring JIRA and GitLab env variables - see README.md for details # from devai.commands.jira import create_jira_issue # from devai.commands.gitlab import create_gitlab_issue_comment -parameters = { - "max_output_tokens": 2048, - "temperature": 0.2 -} - @click.command(name='code') @click.option('-c', '--context', required=False, type=str, default="") def code(context): @@ -58,8 +57,9 @@ def code(context): # Load files as text into source variable source=source.format(format_files_as_string(context)) - code_chat_model = GenerativeModel("gemini-1.0-pro") - code_chat = code_chat_model.start_chat() + code_chat_model = GenerativeModel(model_name) + with telemetry.tool_context_manager(USER_AGENT): + code_chat = code_chat_model.start_chat() code_chat.send_message(qry) response = code_chat.send_message(source) @@ -104,8 +104,9 @@ def performance(context): # Load files as text into source variable source=source.format(format_files_as_string(context)) - code_chat_model = GenerativeModel("gemini-1.0-pro") - code_chat = code_chat_model.start_chat() + code_chat_model = GenerativeModel(model_name) + with telemetry.tool_context_manager(USER_AGENT): + code_chat = code_chat_model.start_chat() code_chat.send_message(qry) response = code_chat.send_message(source) @@ -139,8 +140,9 @@ def security(context): # Load files as text into source variable source=source.format(format_files_as_string(context)) - code_chat_model = GenerativeModel("gemini-1.0-pro") - code_chat = code_chat_model.start_chat() + code_chat_model = GenerativeModel(model_name) + with telemetry.tool_context_manager(USER_AGENT): + code_chat = code_chat_model.start_chat() code_chat.send_message(qry) response = code_chat.send_message(source) diff --git a/outer-loop-cli/src/requirements.txt b/outer-loop-cli/src/requirements.txt index 595a180..6ebeb36 100644 --- a/outer-loop-cli/src/requirements.txt +++ b/outer-loop-cli/src/requirements.txt @@ -1,5 +1,5 @@ click==8.1.7 -google-cloud-aiplatform==1.42.1 +google-cloud-aiplatform==1.48.0 jira==3.6.0 python-gitlab==4.4.0 langchain==0.1.8