diff --git a/README.md b/README.md index b33ba25..7abe474 100644 --- a/README.md +++ b/README.md @@ -124,7 +124,37 @@ In addition to the `OPENAI_API_KEY`, `llm_model` and `model_name` can be edited } ``` -* GCP +* Google GenAI + +Get your Gemini API key via https://aistudio.google.com/app/apikey. + +```json +{ + "llm_config": { + "embedding_service": { + "embedding_model_service": "genai", + "model_name": "models/gemini-embedding-exp-03-07", + "dimensions": 1536, + "authentication_configuration": { + "GOOGLE_API_KEY": "YOUR_GOOGLE_API_KEY_HERE" + } + }, + "completion_service": { + "llm_service": "genai", + "llm_model": "gemini-2.5-flash", + "authentication_configuration": { + "GOOGLE_API_KEY": "YOUR_GOOGLE_API_KEY_HERE" + }, + "model_kwargs": { + "temperature": 0 + }, + "prompt_path": "./common/prompts/google_gemini/" + } + } +} +``` + +* GCP VertexAI Follow the GCP authentication information found here: https://cloud.google.com/docs/authentication/application-default-credentials#GAC and create a Service Account with VertexAI credentials. Then add the following to the docker run command: diff --git a/common/config.py b/common/config.py index 32fb24d..3fed4af 100644 --- a/common/config.py +++ b/common/config.py @@ -9,6 +9,7 @@ AzureOpenAI_Ada002, OpenAI_Embedding, VertexAI_PaLM_Embedding, + GenAI_Embedding, ) from common.embeddings.tigergraph_embedding_store import TigerGraphEmbeddingStore from common.llm_services import ( @@ -16,6 +17,7 @@ AWSBedrock, AzureOpenAI, GoogleVertexAI, + GoogleGenAI, Groq, HuggingFaceEndpoint, LLM_Model, @@ -62,12 +64,20 @@ graphrag_config = server_config.get("graphrag_config") if db_config is None: - raise Exception("graphrag_config is not found in SERVER_CONFIG") + raise Exception("db_config is not found in SERVER_CONFIG") if llm_config is None: - raise Exception("graphrag_config is not found in SERVER_CONFIG") + raise Exception("llm_config is not found in SERVER_CONFIG") + +completion_config = llm_config.get("completion_service") +if completion_config is None: + raise Exception("completion_service is not found in llm_config") +embedding_config = llm_config.get("embedding_service") +if embedding_config is None: + raise Exception("embedding_service is not found in llm_config") +embedding_dimension = embedding_config.get("dimensions", 1536) if graphrag_config is None: - graphrag_config = {"reuse_embedding", true} + graphrag_config = {"reuse_embedding", True} if "chunker" not in graphrag_config: graphrag_config["chunker"] = "semantic" if "extractor" not in graphrag_config: @@ -85,6 +95,8 @@ embedding_service = AzureOpenAI_Ada002(llm_config["embedding_service"]) elif llm_config["embedding_service"]["embedding_model_service"].lower() == "vertexai": embedding_service = VertexAI_PaLM_Embedding(llm_config["embedding_service"]) +elif llm_config["embedding_service"]["embedding_model_service"].lower() == "genai": + embedding_service = GenAI_Embedding(llm_config["embedding_service"]) elif llm_config["embedding_service"]["embedding_model_service"].lower() == "bedrock": embedding_service = AWS_Bedrock_Embedding(llm_config["embedding_service"]) else: @@ -99,6 +111,8 @@ def get_llm_service(llm_config) -> LLM_Model: return AWS_SageMaker_Endpoint(llm_config["completion_service"]) elif llm_config["completion_service"]["llm_service"].lower() == "vertexai": return GoogleVertexAI(llm_config["completion_service"]) + elif llm_config["completion_service"]["llm_service"].lower() == "genai": + return GoogleGenAI(llm_config["completion_service"]) elif llm_config["completion_service"]["llm_service"].lower() == "bedrock": return AWSBedrock(llm_config["completion_service"]) elif llm_config["completion_service"]["llm_service"].lower() == "groq": @@ -129,4 +143,3 @@ def get_llm_service(llm_config) -> LLM_Model: support_ai_instance=True, ) service_status["embedding_store"] = {"status": "ok", "error": None} - diff --git a/common/embeddings/embedding_services.py b/common/embeddings/embedding_services.py index a003793..9c297e9 100644 --- a/common/embeddings/embedding_services.py +++ b/common/embeddings/embedding_services.py @@ -4,6 +4,8 @@ from typing import List from langchain.schema.embeddings import Embeddings +from langchain_community.embeddings.openai import OpenAIEmbeddings +from langchain_google_genai import GoogleGenerativeAIEmbeddings from common.logs.log import req_id_cv from common.logs.logwriter import LogWriter @@ -17,7 +19,7 @@ class EmbeddingModel(Embeddings): Implements connections to the desired embedding API. """ - def __init__(self, config: dict, model_name: str, base_url: str = None): + def __init__(self, config: dict, model_name: str): """Initialize an EmbeddingModel Read JSON config file and export the details as environment variables. """ @@ -27,9 +29,8 @@ def __init__(self, config: dict, model_name: str, base_url: str = None): ] self.embeddings = None self.model_name = model_name - self.base_url = config.get("base_url") LogWriter.info( - f"request_id={req_id_cv.get()} instantiated OpenAI model_name={model_name}" + f"request_id={req_id_cv.get()} instantiated AI model_name={model_name}" ) def embed_documents(self, texts: List[str]) -> List[List[float]]: @@ -47,7 +48,10 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]: try: LogWriter.info(f"request_id={req_id_cv.get()} ENTRY embed_documents()") - docs = self.embeddings.embed_documents(texts) + if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings): + docs = self.embeddings.embed_documents(texts, output_dimensionality=self.dimensions) + else: + docs = self.embeddings.embed_documents(texts) LogWriter.info(f"request_id={req_id_cv.get()} EXIT embed_documents()") metrics.llm_success_response_total.labels(self.model_name).inc() return docs @@ -78,7 +82,10 @@ def embed_query(self, question: str) -> List[float]: logger.debug_pii( f"request_id={req_id_cv.get()} embed_query() embedding question={question}" ) - query_embedding = self.embeddings.embed_query(question) + if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings): + query_embedding = self.embeddings.embed_query(question, output_dimensionality=self.dimensions) + else: + query_embedding = self.embeddings.embed_query(question) LogWriter.info(f"request_id={req_id_cv.get()} EXIT embed_query()") metrics.llm_success_response_total.labels(self.model_name).inc() return query_embedding @@ -105,8 +112,13 @@ async def aembed_query(self, question: str) -> List[float]: # metrics.llm_inprogress_requests.labels(self.model_name).inc() # try: + LogWriter.info(f"request_id={req_id_cv.get()} ENTRY aembed_query()") logger.debug_pii(f"aembed_query() embedding question={question}") - query_embedding = await self.embeddings.aembed_query(question) + if isinstance(self.embeddings, GoogleGenerativeAIEmbeddings): + query_embedding = await self.embeddings.aembed_query(question, output_dimensionality=self.dimensions) + else: + query_embedding = await self.embeddings.aembed_query(question) + LogWriter.info(f"request_id={req_id_cv.get()} EXIT aembed_query()") # metrics.llm_success_response_total.labels(self.model_name).inc() return query_embedding # except Exception as e: @@ -138,10 +150,8 @@ def __init__(self, config): super().__init__( config, model_name=config.get("model_name", "text-embedding-3-small") ) - # from langchain_openai import OpenAIEmbeddings - from langchain_community.embeddings.openai import OpenAIEmbeddings - self.embeddings = OpenAIEmbeddings(model=self.model_name, base_url=self.base_url) + self.embeddings = OpenAIEmbeddings(model=self.model_name, base_url=config.get("base_url")) class VertexAI_PaLM_Embedding(EmbeddingModel): @@ -151,7 +161,18 @@ def __init__(self, config): super().__init__(config, model_name=config.get("model_name", "VertexAI PaLM")) from langchain.embeddings import VertexAIEmbeddings - self.embeddings = VertexAIEmbeddings() + self.embeddings = VertexAIEmbeddings(model_name=self.model_name) + + +class GenAI_Embedding(EmbeddingModel): + """Google GenAI Embedding Model""" + + def __init__(self, config): + super().__init__(config, model_name=config.get("model_name", "gemini-embedding-exp-03-07")) + + self.embeddings = GoogleGenerativeAIEmbeddings(model=self.model_name) + self.dimensions = config.get("dimensions", 1536) + class AWS_Bedrock_Embedding(EmbeddingModel): diff --git a/common/embeddings/tigergraph_embedding_store.py b/common/embeddings/tigergraph_embedding_store.py index b5985c6..c7470ae 100644 --- a/common/embeddings/tigergraph_embedding_store.py +++ b/common/embeddings/tigergraph_embedding_store.py @@ -277,7 +277,7 @@ def remove_embeddings( return def retrieve_similar(self, query_embedding, top_k=10, filter_expr: str = None, vertex_types: List[str] = ["DocumentChunk"]): - res = retrieve_similar_with_score(query_embedding, top_k=top_k, filter_expr=filter_expr, vertex_types=vertex_types) + res = self.retrieve_similar_with_score(query_embedding, top_k=top_k, filter_expr=filter_expr, vertex_types=vertex_types) similar = [x[0] for x in res] return similar diff --git a/common/llm_services/__init__.py b/common/llm_services/__init__.py index 3fde87f..18dc49e 100644 --- a/common/llm_services/__init__.py +++ b/common/llm_services/__init__.py @@ -3,8 +3,9 @@ from .openai_service import OpenAI from .aws_sagemaker_endpoint import AWS_SageMaker_Endpoint from .google_vertexai_service import GoogleVertexAI +from .google_genai_service import GoogleGenAI from .aws_bedrock_service import AWSBedrock from .groq_llm_service import Groq from .ollama import Ollama from .huggingface_endpoint import HuggingFaceEndpoint -from .ibm_watsonx_service import IBMWatsonX \ No newline at end of file +from .ibm_watsonx_service import IBMWatsonX diff --git a/common/llm_services/google_genai_service.py b/common/llm_services/google_genai_service.py new file mode 100644 index 0000000..a96716d --- /dev/null +++ b/common/llm_services/google_genai_service.py @@ -0,0 +1,114 @@ +import logging +import os + +from common.llm_services import LLM_Model +from langchain_google_genai import ChatGoogleGenerativeAI + +from common.logs.log import req_id_cv +from common.logs.logwriter import LogWriter + +logger = logging.getLogger(__name__) + + +class GoogleGenAI(LLM_Model): + def __init__(self, config): + super().__init__(config) + for auth_detail in config["authentication_configuration"].keys(): + os.environ[auth_detail] = config["authentication_configuration"][ + auth_detail + ] + + model_name = config["llm_model"] + self.llm = ChatGoogleGenerativeAI( + temperature=config["model_kwargs"]["temperature"], + model=model_name, + max_tokens=None, + timeout=None, + max_retries=2, + ) + self.prompt_path = config["prompt_path"] + LogWriter.info( + f"request_id={req_id_cv.get()} instantiated OpenAI model_name={model_name}" + ) + + @property + def map_question_schema_prompt(self): + return self._read_prompt_file(self.prompt_path + "map_question_to_schema.txt") + + @property + def generate_function_prompt(self): + return self._read_prompt_file(self.prompt_path + "generate_function.txt") + + @property + def generate_cypher_prompt(self): + filepath = self.prompt_path + "generate_cypher.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().generate_cypher_prompt + + @property + def entity_relationship_extraction_prompt(self): + return self._read_prompt_file( + self.prompt_path + "entity_relationship_extraction.txt" + ) + + @property + def route_response_prompt(self): + filepath = self.prompt_path + "route_response.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().route_response_prompt + + @property + def graphrag_scoring_prompt(self): + filepath = self.prompt_path + "graphrag_scoring.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().graphrag_scoring_prompt + + @property + def keyword_extraction_prompt(self): + filepath = self.prompt_path + "keyword_extraction.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().keyword_extraction_prompt + + @property + def question_expansion_prompt(self): + filepath = self.prompt_path + "question_expansion.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().question_expansion_prompt + + @property + def supportai_response_prompt(self): + filepath = self.prompt_path + "supportai_response.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().supportai_response_prompt + + @property + def chatbot_response_prompt(self): + filepath = self.prompt_path + "chatbot_response.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().chatbot_response_prompt + + @property + def hyde_prompt(self): + filepath = self.prompt_path + "hyde.txt" + if os.path.exists(filepath): + return self._read_prompt_file(filepath) + else: + return super().hyde_prompt + + @property + def model(self): + return self.llm diff --git a/common/llm_services/openai_service.py b/common/llm_services/openai_service.py index e85ada8..6fa5d75 100644 --- a/common/llm_services/openai_service.py +++ b/common/llm_services/openai_service.py @@ -1,10 +1,7 @@ import logging import os -if os.getenv("ECC"): - from langchain_openai.chat_models import ChatOpenAI -else: - from langchain_community.chat_models import ChatOpenAI +from langchain_openai.chat_models import ChatOpenAI from common.llm_services import LLM_Model from common.logs.log import req_id_cv diff --git a/common/prompts/google_gemini/chatbot_response.txt b/common/prompts/google_gemini/chatbot_response.txt new file mode 100644 index 0000000..c412948 --- /dev/null +++ b/common/prompts/google_gemini/chatbot_response.txt @@ -0,0 +1,32 @@ +You are a highly efficient and empathetic AI-powered customer support agent for TigerGraph. Your goal is to provide accurate, helpful, and friendly support while maintaining professionalism. Follow these guidelines: + +1. **Understand the User's Needs**: Carefully interpret the user's question or issue. Ask clarifying questions if necessary. +2. **Provide Accurate Answers**: Use the given context in JSON format to deliver reliable and precise responses. +3. **Be Polite and Friendly**: Maintain a tone that is warm, respectful, and supportive. +4. **Stay Within Scope**: Respond only with information relevant to the user's question. If the information is unavailable, apologize and suggest alternative actions. +5. **Encourage Next Steps**: Proactively suggest solutions or next steps, such as contacting a human agent if the issue cannot be resolved by the AI. +6. **Escalate When Needed**: Recognize complex or sensitive queries that require human attention and politely guide the user to contact human support. +7. **Reference Relevant Knowledge**: Ground your responses based on the context and resources. Include the IDs of the knowledge (always in UUID format) you refer to in your response in the format [UUID]. +8. **Here are the channels to reach the TigerGraph team**: + - [Contact Sales](https://www.tigergraph.com/contact/) + - [Support Portal](https://www.tigergraph.com/support/) + - [Documentation](https://docs.tigergraph.com/) + - [Community Forum](https://dev.tigergraph.com/) + - [Email](mailto:support@tigergraph.com) + +Example format for responses: +- **Solution**: "Here's the information you requested: [details]." +- **Clarification**: "Could you provide more details about [specific aspect]?" +- **Escalation**: "This issue may require human assistance. Please reach out via [Support Portal](https://www.tigergraph.com/support/) or email us at [support@tigergraph.com](mailto:support@tigergraph.com)." + +Your mission is to ensure a seamless, satisfying customer experience while upholding TigerGraph's values and commitment to enterprise excellence. + +Give the context in JSON format, combine and rephrase it to answer the question. +Use only the provided information in context without adding any reasoning or additional logic. +Make sure all information in the context are covered in the generated answer. +Generate the answer in JSON format, make sure to escape necessary characters in order to return a valid JSON response only. +Make sure all the fields required by the format instructions are included, set a field to empty if you don't have that information. + +Question: {question} +Context: {context} +Format: {format_instructions} diff --git a/common/prompts/google_gemini/entity_relationship_extraction.txt b/common/prompts/google_gemini/entity_relationship_extraction.txt new file mode 100644 index 0000000..852dded --- /dev/null +++ b/common/prompts/google_gemini/entity_relationship_extraction.txt @@ -0,0 +1,24 @@ +# Knowledge Graph Instructions for GPT-4 +## 1. Overview +You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph. +- **Nodes** represent entities, concepts, and properties of entities. +- The aim is to achieve simplicity and clarity in the knowledge graph, making it accessible for a vast audience. +## 2. Labeling Nodes +- **Consistency**: Ensure you use basic or elementary types for node labels. +- For example, when you identify an entity representing a person, always label it as **"person"**. Avoid using more specific terms like "mathematician" or "scientist". +- **Node IDs**: Never utilize integers as node IDs. Node IDs should be names or human-readable identifiers found in the text. +## 3. Handling Numerical Data and Dates +- Numerical data, like age or other related information, should be incorporated as attributes or properties of the respective nodes. +- **No Separate Nodes for Dates/Numbers**: Do not create separate nodes for dates or numerical values. Always attach them as attributes or properties of nodes. +- **Property Format**: Properties must be in a key-value format. Only use properties for dates and numbers, string properties should be new nodes. +- **Quotation Marks**: Never use escaped single or double quotes within property values. +- **Naming Convention**: Use camelCase for property keys, e.g., `birthDate`. +## 4. Coreference Resolution +- **Maintain Entity Consistency**: When extracting entities, it's vital to ensure consistency. +If an entity, such as "John Doe", is mentioned multiple times in the text but is referred to by different names or pronouns (e.g., "Joe", "he"), +always use the most complete identifier for that entity throughout the knowledge graph. In this example, use "John Doe" as the entity ID. +Remember, the knowledge graph should be coherent and easily understandable, so maintaining consistency in entity references is crucial. +## 5. Strict Compliance +Adhere to the rules strictly. Non-compliance will result in termination, including poor formatting. +## 6. Handling Instances with No Relationships +If a node has no relationships, it should still be included in the knowledge graph. Simply add the node and leave the relationships section empty. \ No newline at end of file diff --git a/common/prompts/google_gemini/generate_cypher.txt b/common/prompts/google_gemini/generate_cypher.txt new file mode 100644 index 0000000..91fcbe6 --- /dev/null +++ b/common/prompts/google_gemini/generate_cypher.txt @@ -0,0 +1,60 @@ +You're an expert in OpenCypher programming. Given the following schema: {schema}, what is the OpenCypher query that retrieves the answer for question: {question}. +If there're multiple words in the question having same meaning then remove the duplication. +Only include attributes that are found in the schema. Never include any attributes that are not found in the schema. +Use attributes instead of primary id if attribute name is more similar to the keyword type in the question. +Use as less vertex type, edge type and attributes as possible. If an attribute is not found in the schema, please exclude it from the query. +Do not return attributes that are not explicitly mentioned in the question. If a vertex type is mentioned in the question, only return the vertex. +Never use directed edge pattern in the OpenCypher query. Always use and create query using undirected pattern. Always ensure the edge used starts from and ends with correct vertex types matching the schema. +Always use double quotes for strings instead of single quotes. +Use alias for ORDER BY if any, and make sure the same alias is used in both RETURN and ORDER BY. Always add ASC or DESC for ORDER BY based on data type. + +Only use the Supported Clauses, Operators, Functions and Expressions below but do not use any of the Unsupported Features, Functions or Syntax Limitations below: + +Supported Clauses: +MATCH / OPTIONAL MATCH / MANDATORY MATCH: Match patterns in the graph. +WHERE: Filter results. +RETURN / WITH: Project query results, alias fields, chain query parts. +ORDER BY / SKIP / LIMIT: Control output order, offset, and size. +DELETE / DETACH DELETE: Delete nodes/edges. + +Supported Operators: +Mathematical: +, -, *, /, %, ^ (exponent) +Comparison: =, <, <=, >, >=, <>, IS NULL, IS NOT NULL +Boolean: AND, OR, NOT, XOR +String/List: CONTAINS, STARTS WITH, ENDS WITH, IN, DISTINCT, [ ] (subscript), . (property access) + +Supported Functions: +Aggregation: count(), sum(), avg(), min(), max(), stDev(), stDevP() +Math: abs(), sqrt(), log(), exp(), sin(), cos(), tan(), radians(), degrees() +String: left(), right(), substring(), replace(), trim(), toLower(), toUpper(), split() +List: head(), last(), size(), range(), coalesce(), tail() +Others: id(), elementId(), labels(), properties(), timestamp() + +Supported Expressions: +CASE: Conditional logic. + +Unsupported Features: +Clauses Not Yet Supported +CALL, CREATE, MERGE, REMOVE, SET, UNION, UNION ALL, UNWIND + +Unsupported Functions: +collect(), exists(), keys(), nodes(), relationships(), length(), percentileCont(), percentileDisc(), startNode(), endNode(), reverse() (list form) + +Syntax Limitations: +WITH clause must group by exactly one vertex variable. +Path variables (e.g. p = (...)) not supported. +MATCH must reference variables from prior WITH. +Disconnected MATCH fragments not supported. + +Additionally, you cannot use the following clauses: +CREATE +MERGE +REMOVE +UNION +UNION ALL +UNWIND +SET + +Make sure to have correct attribute names in the OpenCypher query and not to name result aliases that are vertex or edge types, operator or function names, and other reserved keywords. + +ONLY write the OpenCypher query in the response. Do not include any other information in the response. diff --git a/common/prompts/google_gemini/generate_function.txt b/common/prompts/google_gemini/generate_function.txt new file mode 100644 index 0000000..781d63d --- /dev/null +++ b/common/prompts/google_gemini/generate_function.txt @@ -0,0 +1,24 @@ +Use the vertex types, edge types, and their attributes and IDs below to write the pyTigerGraph function call to answer the question using a pyTigerGraph connection. +When the question asks for "How many", make sure to always select a function that contains "Count" in the description/function call. Make sure never to generate a function that is not listed below. +When certain entities are mapped to vertex attributes, may consider to generate a WHERE clause. +If a WHERE clause is generated, please follow the instruction with proper quoting. To construct a WHERE clause string. Ensure that string attribute values are properly quoted. +For example, if the generated function contains "('Person', where='name=William Torres')", Expected Output: "('Person', where='name="William Torres"')", This rule applies to all types of attributes. e.g., name, email, address and so on. +Documentation contains helpful Python docstrings for the various functions. Use this knowledge to construct the proper function call. Choose one function to execute. +Don't generate target_vertex_ids if there is no the term 'id' explicitly mentioned in the question. +Vertex Types: {vertex_types} +Vertex Attributes: {vertex_attributes} +Vertex IDs: {vertex_ids} +Edge Types: {edge_types} +Edge Attributes: {edge_attributes} +Question: {question} +First Docstring: {doc1} +Second Docstring: {doc2} +Third Docstring: {doc3} +Fourth Docstring: {doc4} +Fifth Docstring: {doc5} +Sixth Docstring: {doc6} +Seventh Docstring: {doc7} +Eighth Docstring: {doc8} + +Follow the output directions below on how to structure your response: +{format_instructions} diff --git a/common/prompts/google_gemini/graphrag_scoring.txt b/common/prompts/google_gemini/graphrag_scoring.txt new file mode 100644 index 0000000..38ef643 --- /dev/null +++ b/common/prompts/google_gemini/graphrag_scoring.txt @@ -0,0 +1,7 @@ +You are a helpful assistant responsible for generating an answer to the question below using the data provided. +Include a quality score for the answer, based on how well it answers the question. The quality score should be between 0 (poor) and 100 (excellent). + +Question: {question} +Context: {context} + +{format_instructions} diff --git a/common/prompts/google_gemini/map_question_to_schema.txt b/common/prompts/google_gemini/map_question_to_schema.txt new file mode 100644 index 0000000..81ed53d --- /dev/null +++ b/common/prompts/google_gemini/map_question_to_schema.txt @@ -0,0 +1,15 @@ +Replace the entites mentioned in the question to one of these choices: {vertices}. +If an entity, such as "John Doe", is mentioned multiple times in the conversation but is referred to by different names or pronouns (e.g., "Joe", "he"), +always use the most complete identifier for that entity throughout the question. In this example, use "John Doe" as the entity. +Choose a better mapping between vertex type or its attributes: {verticesAttrs}. +Replace the relationships mentioned in the question to one of these choices: {edges}. +Make sure the entities are either the source vertices or target vertices of the relationships: {edgesInfo}. +When certain entities are mapped to vertex attributes, may consider to generate a WHERE clause. +If there are words that are synonyms with the entities or relationships above, make sure to output the cannonical form found in the choices above. +Generate the complete question with the appropriate replacements. Keep the case of the schema elements the same. +Don't generate target_vertex_ids if there is no the term 'id' explicitly mentioned in the question. + +{format_instructions} +question: {question} +conversation: {conversation} + diff --git a/common/prompts/google_gemini/question_expansion.txt b/common/prompts/google_gemini/question_expansion.txt new file mode 100644 index 0000000..b04aed8 --- /dev/null +++ b/common/prompts/google_gemini/question_expansion.txt @@ -0,0 +1,6 @@ +You are a helpful assistant responsible for generating 10 new questions similar to the original question below to represent its meaning in a more clear way. +Include a quality score for the answer, based on how well it represents the meaning of the original question. The quality score should be between 0 (poor) and 100 (excellent). + +Question: {question} + +{format_instructions} diff --git a/common/prompts/openai_gpt4/chatbot_response.txt b/common/prompts/openai_gpt4/chatbot_response.txt new file mode 100644 index 0000000..c412948 --- /dev/null +++ b/common/prompts/openai_gpt4/chatbot_response.txt @@ -0,0 +1,32 @@ +You are a highly efficient and empathetic AI-powered customer support agent for TigerGraph. Your goal is to provide accurate, helpful, and friendly support while maintaining professionalism. Follow these guidelines: + +1. **Understand the User's Needs**: Carefully interpret the user's question or issue. Ask clarifying questions if necessary. +2. **Provide Accurate Answers**: Use the given context in JSON format to deliver reliable and precise responses. +3. **Be Polite and Friendly**: Maintain a tone that is warm, respectful, and supportive. +4. **Stay Within Scope**: Respond only with information relevant to the user's question. If the information is unavailable, apologize and suggest alternative actions. +5. **Encourage Next Steps**: Proactively suggest solutions or next steps, such as contacting a human agent if the issue cannot be resolved by the AI. +6. **Escalate When Needed**: Recognize complex or sensitive queries that require human attention and politely guide the user to contact human support. +7. **Reference Relevant Knowledge**: Ground your responses based on the context and resources. Include the IDs of the knowledge (always in UUID format) you refer to in your response in the format [UUID]. +8. **Here are the channels to reach the TigerGraph team**: + - [Contact Sales](https://www.tigergraph.com/contact/) + - [Support Portal](https://www.tigergraph.com/support/) + - [Documentation](https://docs.tigergraph.com/) + - [Community Forum](https://dev.tigergraph.com/) + - [Email](mailto:support@tigergraph.com) + +Example format for responses: +- **Solution**: "Here's the information you requested: [details]." +- **Clarification**: "Could you provide more details about [specific aspect]?" +- **Escalation**: "This issue may require human assistance. Please reach out via [Support Portal](https://www.tigergraph.com/support/) or email us at [support@tigergraph.com](mailto:support@tigergraph.com)." + +Your mission is to ensure a seamless, satisfying customer experience while upholding TigerGraph's values and commitment to enterprise excellence. + +Give the context in JSON format, combine and rephrase it to answer the question. +Use only the provided information in context without adding any reasoning or additional logic. +Make sure all information in the context are covered in the generated answer. +Generate the answer in JSON format, make sure to escape necessary characters in order to return a valid JSON response only. +Make sure all the fields required by the format instructions are included, set a field to empty if you don't have that information. + +Question: {question} +Context: {context} +Format: {format_instructions} diff --git a/common/prompts/openai_gpt4/generate_cypher.txt b/common/prompts/openai_gpt4/generate_cypher.txt new file mode 100644 index 0000000..91fcbe6 --- /dev/null +++ b/common/prompts/openai_gpt4/generate_cypher.txt @@ -0,0 +1,60 @@ +You're an expert in OpenCypher programming. Given the following schema: {schema}, what is the OpenCypher query that retrieves the answer for question: {question}. +If there're multiple words in the question having same meaning then remove the duplication. +Only include attributes that are found in the schema. Never include any attributes that are not found in the schema. +Use attributes instead of primary id if attribute name is more similar to the keyword type in the question. +Use as less vertex type, edge type and attributes as possible. If an attribute is not found in the schema, please exclude it from the query. +Do not return attributes that are not explicitly mentioned in the question. If a vertex type is mentioned in the question, only return the vertex. +Never use directed edge pattern in the OpenCypher query. Always use and create query using undirected pattern. Always ensure the edge used starts from and ends with correct vertex types matching the schema. +Always use double quotes for strings instead of single quotes. +Use alias for ORDER BY if any, and make sure the same alias is used in both RETURN and ORDER BY. Always add ASC or DESC for ORDER BY based on data type. + +Only use the Supported Clauses, Operators, Functions and Expressions below but do not use any of the Unsupported Features, Functions or Syntax Limitations below: + +Supported Clauses: +MATCH / OPTIONAL MATCH / MANDATORY MATCH: Match patterns in the graph. +WHERE: Filter results. +RETURN / WITH: Project query results, alias fields, chain query parts. +ORDER BY / SKIP / LIMIT: Control output order, offset, and size. +DELETE / DETACH DELETE: Delete nodes/edges. + +Supported Operators: +Mathematical: +, -, *, /, %, ^ (exponent) +Comparison: =, <, <=, >, >=, <>, IS NULL, IS NOT NULL +Boolean: AND, OR, NOT, XOR +String/List: CONTAINS, STARTS WITH, ENDS WITH, IN, DISTINCT, [ ] (subscript), . (property access) + +Supported Functions: +Aggregation: count(), sum(), avg(), min(), max(), stDev(), stDevP() +Math: abs(), sqrt(), log(), exp(), sin(), cos(), tan(), radians(), degrees() +String: left(), right(), substring(), replace(), trim(), toLower(), toUpper(), split() +List: head(), last(), size(), range(), coalesce(), tail() +Others: id(), elementId(), labels(), properties(), timestamp() + +Supported Expressions: +CASE: Conditional logic. + +Unsupported Features: +Clauses Not Yet Supported +CALL, CREATE, MERGE, REMOVE, SET, UNION, UNION ALL, UNWIND + +Unsupported Functions: +collect(), exists(), keys(), nodes(), relationships(), length(), percentileCont(), percentileDisc(), startNode(), endNode(), reverse() (list form) + +Syntax Limitations: +WITH clause must group by exactly one vertex variable. +Path variables (e.g. p = (...)) not supported. +MATCH must reference variables from prior WITH. +Disconnected MATCH fragments not supported. + +Additionally, you cannot use the following clauses: +CREATE +MERGE +REMOVE +UNION +UNION ALL +UNWIND +SET + +Make sure to have correct attribute names in the OpenCypher query and not to name result aliases that are vertex or edge types, operator or function names, and other reserved keywords. + +ONLY write the OpenCypher query in the response. Do not include any other information in the response. diff --git a/common/py_schemas/tool_io_schemas.py b/common/py_schemas/tool_io_schemas.py index fa81355..af81038 100644 --- a/common/py_schemas/tool_io_schemas.py +++ b/common/py_schemas/tool_io_schemas.py @@ -1,6 +1,6 @@ from typing import Dict, List, Optional -from langchain.pydantic_v1 import BaseModel, Field +from pydantic import BaseModel, Field from langchain_community.graphs.graph_document import Node as BaseNode from langchain_community.graphs.graph_document import Relationship as BaseRelationship diff --git a/common/requirements.txt b/common/requirements.txt index 03df9f1..e8f791a 100644 --- a/common/requirements.txt +++ b/common/requirements.txt @@ -1,174 +1,175 @@ -aiochannel==1.2.1 -aiohappyeyeballs==2.3.5 -aiohttp==3.10.3 -aiosignal==1.3.1 +aiochannel==1.3.0 +aiohappyeyeballs==2.6.1 +aiohttp==3.12.13 +aiosignal==1.3.2 annotated-types==0.7.0 -anyio==4.4.0 +anyio==4.9.0 appdirs==1.4.4 -argon2-cffi==23.1.0 +argon2-cffi==25.1.0 argon2-cffi-bindings==21.2.0 -async-timeout==4.0.3 -asyncer==0.0.7 -attrs==24.2.0 -azure-core==1.30.2 -azure-storage-blob==12.22.0 +async-timeout==5.0.1 +asyncer==0.0.8 +attrs==25.3.0 +azure-core==1.34.0 +azure-storage-blob==12.25.1 backoff==2.2.1 -beautifulsoup4==4.12.3 -boto3==1.34.159 -botocore==1.34.159 -cachetools==5.4.0 -certifi==2024.7.4 -cffi==1.17.0 +beautifulsoup4==4.13.4 +boto3==1.38.45 +botocore==1.38.45 +cachetools==5.5.2 +certifi==2025.6.15 +cffi==1.17.1 chardet==5.2.0 -charset-normalizer==3.3.2 -click==8.1.7 -contourpy==1.2.1 -cryptography==43.0.0 +charset-normalizer==3.4.2 +click==8.2.1 +contourpy==1.3.2 +cryptography==45.0.4 cycler==0.12.1 dataclasses-json==0.6.7 -deepdiff==7.0.1 +deepdiff==8.5.0 distro==1.9.0 docker-pycreds==0.4.0 docstring_parser==0.16 -emoji==2.12.1 -environs==9.5.0 -exceptiongroup==1.2.2 -fastapi==0.112.0 -filelock==3.15.4 +emoji==2.14.1 +environs==14.2.0 +exceptiongroup==1.3.0 +fastapi==0.115.14 +filelock==3.18.0 filetype==1.2.0 -fonttools==4.53.1 -frozenlist==1.4.1 -fsspec==2024.6.1 -gitdb==4.0.11 -GitPython==3.1.43 -google-api-core==2.19.1 -google-auth==2.33.0 -google-cloud-aiplatform==1.61.0 -google-cloud-bigquery==3.25.0 -google-cloud-core==2.4.1 -google-cloud-resource-manager==1.12.5 -google-cloud-storage==2.18.2 -google-crc32c==1.5.0 +fonttools==4.58.4 +frozenlist==1.7.0 +fsspec==2025.5.1 +gitdb==4.0.12 +GitPython==3.1.44 +google-api-core==2.25.1 +google-auth==2.40.3 +google-cloud-aiplatform==1.99.0 +google-cloud-bigquery==3.34.0 +google-cloud-core==2.4.3 +google-cloud-resource-manager==1.14.2 +google-cloud-storage==2.19.0 +google-crc32c==1.7.1 google-resumable-media==2.7.2 -googleapis-common-protos==1.63.2 -greenlet==3.0.3 -groq==0.9.0 -grpc-google-iam-v1==0.13.1 -grpcio==1.63.0 -grpcio-status==1.63.0 -h11==0.14.0 -httpcore==1.0.5 -httptools==0.6.1 -httpx==0.27.0 -huggingface-hub==0.24.5 -ibm-cos-sdk==2.13.6 -ibm-cos-sdk-core==2.13.6 -ibm-cos-sdk-s3transfer==2.13.6 -ibm_watsonx_ai==1.1.5 -idna==3.7 -importlib_metadata==8.2.0 -iniconfig==2.0.0 -isodate==0.6.1 -jiter==0.5.0 +googleapis-common-protos==1.70.0 +greenlet==3.2.3 +groq==0.29.0 +grpc-google-iam-v1==0.14.2 +grpcio==1.73.1 +grpcio-status==1.73.1 +h11==0.16.0 +httpcore==1.0.9 +httptools==0.6.4 +httpx==0.28.1 +huggingface-hub==0.33.1 +ibm-cos-sdk==2.14.2 +ibm-cos-sdk-core==2.14.2 +ibm-cos-sdk-s3transfer==2.14.2 +ibm_watsonx_ai==1.3.26 +idna==3.10 +importlib_metadata==8.7.0 +iniconfig==2.1.0 +isodate==0.7.2 +jiter==0.10.0 jmespath==1.0.1 -joblib==1.4.2 -jq==1.7.0 +joblib==1.5.1 +jq==1.9.1 jsonpatch==1.33 jsonpath-python==1.0.6 jsonpointer==3.0.0 -kiwisolver==1.4.5 -langchain==0.2.13 -langchain-community==0.2.12 -langchain-core==0.2.30 -langchain-experimental==0.0.64 -langchain-groq==0.1.9 -langchain-ibm==0.1.12 -langchain-openai==0.1.21 -langchain-text-splitters==0.2.2 +kiwisolver==1.4.8 +langchain==0.3.26 +langchain-core==0.3.66 +langchain_google_genai==2.1.5 +langchain-community==0.3.26 +langchain-experimental==0.3.5rc1 +langchain-groq==0.3.4 +langchain-ibm==0.3.12 +langchain-openai==0.3.26 +langchain-text-splitters==0.3.8 langchainhub==0.1.21 langdetect==1.0.9 -langgraph==0.2.3 -langgraph-checkpoint==1.0.2 -langsmith==0.1.99 -Levenshtein==0.25.1 +langgraph==0.4.10 +langgraph-checkpoint==2.1.0 +langsmith==0.4.2 +Levenshtein==0.27.1 lomond==0.3.3 -lxml==5.3.0 -marshmallow==3.21.3 -matplotlib==3.9.2 -multidict==6.0.5 -mypy-extensions==1.0.0 +lxml==6.0.0 +marshmallow==3.26.1 +matplotlib==3.10.3 +multidict==6.5.1 +mypy-extensions==1.1.0 nest-asyncio==1.6.0 -nltk==3.8.1 -numpy==1.26.4 -openai==1.40.6 +nltk==3.9.1 +numpy==2.3.1 +openai==1.92.2 ordered-set==4.1.0 -orjson==3.10.7 -packaging==24.1 -pandas==2.1.4 +orjson==3.10.18 +packaging==24.2 +pandas==2.2.3 pathtools==0.1.2 -pillow==10.4.0 -platformdirs==4.2.2 -pluggy==1.5.0 -prometheus_client==0.20.0 -proto-plus==1.24.0 -protobuf==5.27.3 -psutil==6.0.0 -pyarrow==17.0.0 -pyasn1==0.6.0 -pyasn1_modules==0.4.0 +pillow==11.2.1 +platformdirs==4.3.8 +pluggy==1.6.0 +prometheus_client==0.22.1 +proto-plus==1.26.1 +protobuf==6.31.1 +psutil==7.0.0 +pyarrow==20.0.0 +pyasn1==0.6.1 +pyasn1_modules==0.4.2 pycparser==2.22 -pycryptodome==3.20.0 -pydantic==2.8.2 -pydantic_core==2.20.1 -pygit2==1.15.1 -pyparsing==3.1.2 -pypdf==4.3.1 -pytest==8.3.2 +pycryptodome==3.23.0 +pydantic==2.11.7 +pydantic_core==2.33.2 +pygit2==1.18.0 +pyparsing==3.2.3 +pypdf==5.6.1 +pytest==8.4.1 python-dateutil==2.9.0.post0 -python-dotenv==1.0.1 -python-iso639==2024.4.27 +python-dotenv==1.1.1 +python-iso639==2025.2.18 python-magic==0.4.27 pyTigerDriver==1.0.15 pyTigerGraph==1.9.0 -pytz==2024.1 +pytz==2025.2 PyYAML==6.0.2 -rapidfuzz==3.9.6 -regex==2024.7.24 -requests==2.32.2 +rapidfuzz==3.13.0 +regex==2024.11.6 +requests==2.32.4 requests-toolbelt==1.0.0 -rsa==4.9 -s3transfer==0.10.2 -scikit-learn==1.5.1 -scipy==1.14.0 -sentry-sdk==2.13.0 -setproctitle==1.3.3 -shapely==2.0.5 -six==1.16.0 -smmap==5.0.1 +rsa==4.9.1 +s3transfer==0.13.0 +scikit-learn==1.7.0 +scipy==1.16.0 +sentry-sdk==2.31.0 +setproctitle==1.3.6 +shapely==2.1.1 +six==1.17.0 +smmap==5.0.2 sniffio==1.3.1 -soupsieve==2.6 -SQLAlchemy==2.0.32 -starlette==0.37.2 +soupsieve==2.7 +SQLAlchemy==2.0.41 +starlette==0.46.2 tabulate==0.9.0 -tenacity==8.5.0 -threadpoolctl==3.5.0 -tiktoken==0.7.0 -tqdm==4.66.5 -types-requests==2.32.0.20240712 +tenacity==9.1.2 +threadpoolctl==3.6.0 +tiktoken==0.9.0 +tqdm==4.67.1 +types-requests==2.32.4.20250611 types-urllib3==1.26.25.14 typing-inspect==0.9.0 -typing_extensions==4.12.2 -tzdata==2024.1 +typing_extensions==4.14.0 +tzdata==2025.2 ujson==5.10.0 -unstructured==0.15.1 -unstructured-client==0.25.5 -urllib3==2.2.2 -uvicorn==0.30.6 -uvloop==0.19.0 -validators==0.33.0 -wandb==0.17.6 -watchfiles==0.23.0 -websockets==12.0 -wrapt==1.16.0 -yarl==1.9.4 -zipp==3.20.0 +unstructured==0.18.1 +unstructured-client==0.37.2 +urllib3==2.5.0 +uvicorn==0.34.3 +uvloop==0.21.0 +validators==0.35.0 +wandb==0.20.1 +watchfiles==1.1.0 +websockets==15.0.1 +wrapt==1.17.2 +yarl==1.20.1 +zipp==3.23.0 diff --git a/docs/notebooks/TransactionFraudInvestigation.ipynb b/docs/notebooks/TransactionFraudInvestigation.ipynb index fc589d8..9e313c3 100644 --- a/docs/notebooks/TransactionFraudInvestigation.ipynb +++ b/docs/notebooks/TransactionFraudInvestigation.ipynb @@ -81,7 +81,7 @@ "from langchain.prompts import PromptTemplate\n", "from langchain_core.output_parsers import PydanticOutputParser\n", "\n", - "from langchain.pydantic_v1 import BaseModel, Field\n", + "from pydantic import BaseModel, Field\n", "\n", "class Questions(BaseModel):\n", " questions: dict[str, str] = Field(description=\"A dictionary of questions to ask the data analyst. The key is the question and the value is the explanation on why to ask the question.\")\n", @@ -423,7 +423,7 @@ "from langchain.prompts import PromptTemplate\n", "from langchain_core.output_parsers import PydanticOutputParser\n", "\n", - "from langchain.pydantic_v1 import BaseModel, Field\n", + "from pydantic import BaseModel, Field\n", "\n", "class DraftReport(BaseModel):\n", " draft_report: str = Field(description=\"A draft of the report based on the answers received from the data analyst. Include citations by adding `[x]` where x is the function call that was used to determine the answer.\")\n", @@ -2078,7 +2078,7 @@ "from langchain.prompts import PromptTemplate\n", "from langchain_core.output_parsers import PydanticOutputParser\n", "\n", - "from langchain.pydantic_v1 import BaseModel, Field\n", + "from pydantic import BaseModel, Field\n", "\n", "class Report(BaseModel):\n", " report: str = Field(description=\"A drareport based on the answers received from the data analyst. Include citations by adding `[x]` where x is the function call that was used to determine the answer.\")\n", diff --git a/docs/notebooks/VisualizeAgent.ipynb b/docs/notebooks/VisualizeAgent.ipynb index 2547133..e85cab0 100644 --- a/docs/notebooks/VisualizeAgent.ipynb +++ b/docs/notebooks/VisualizeAgent.ipynb @@ -40,8 +40,8 @@ " TigerGraphConnection(),\n", " EmbeddingModel({\"authentication_configuration\": {}}, \"\"),\n", " EmbeddingStore,\n", - " MapQuestionToSchema(TigerGraphConnection(), LLM_Model({}), \"\"),\n", - " GenerateFunction(TigerGraphConnection(), LLM_Model({}), \"\", EmbeddingModel({\"authentication_configuration\": {}}, \"\"), EmbeddingStore),\n", + " MapQuestionToSchema(TigerGraphConnection(), LLM_Model({})),\n", + " GenerateFunction(TigerGraphConnection(), LLM_Model({}), EmbeddingModel({\"authentication_configuration\": {}}, \"\"), EmbeddingStore),\n", " GenerateCypher(TigerGraphConnection(), LLM_Model({}))\n", ").create_graph()" ] diff --git a/ecc/app/ecc_util.py b/ecc/app/ecc_util.py index 8509314..7c28507 100644 --- a/ecc/app/ecc_util.py +++ b/ecc/app/ecc_util.py @@ -5,6 +5,7 @@ AWSBedrock, AzureOpenAI, GoogleVertexAI, + GoogleGenAI, Groq, HuggingFaceEndpoint, Ollama, @@ -50,6 +51,8 @@ def get_llm_service(): llm_provider = AWS_SageMaker_Endpoint(llm_config["completion_service"]) elif llm_config["completion_service"]["llm_service"].lower() == "vertexai": llm_provider = GoogleVertexAI(llm_config["completion_service"]) + elif llm_config["completion_service"]["llm_service"].lower() == "genai": + llm_provider = GoogleGenAI(llm_config["completion_service"]) elif llm_config["completion_service"]["llm_service"].lower() == "bedrock": llm_provider = AWSBedrock(llm_config["completion_service"]) elif llm_config["completion_service"]["llm_service"].lower() == "groq": diff --git a/ecc/app/graphrag/graph_rag.py b/ecc/app/graphrag/graph_rag.py index 3dd94fc..1e1f81c 100644 --- a/ecc/app/graphrag/graph_rag.py +++ b/ecc/app/graphrag/graph_rag.py @@ -79,7 +79,7 @@ async def chunk_docs( Creates and starts one worker for each document in the docs channel. """ - logger.info("Reading from docs channel") + logger.info("Chunk Processing Start") doc_tasks = [] async with asyncio.TaskGroup() as grp: while True: @@ -94,7 +94,7 @@ async def chunk_docs( except Exception: raise - logger.info("chunk_docs done") + logger.info("Chunk Processing End") # close the extract chan -- chunk_doc is the only sender # and chunk_doc calls are kicked off from here @@ -109,7 +109,7 @@ async def upsert(upsert_chan: Channel): (func, args) <- q.get() """ - logger.info("Reading from upsert channel") + logger.info("Data Upserting Start") # consume task queue async with asyncio.TaskGroup() as grp: while True: @@ -123,13 +123,13 @@ async def upsert(upsert_chan: Channel): except Exception: raise - logger.info("upsert done") + logger.info("Data Upserting End") logger.info("closing load_q chan") load_q.close() async def load(conn: AsyncTigerGraphConnection): - logger.info("Reading from load_q") + logger.info("Data Loading Start") dd = lambda: defaultdict(dd) # infinite default dict batch_size = 500 # while the load q is still open or has contents @@ -190,6 +190,7 @@ async def load(conn: AsyncTigerGraphConnection): # TODO: flush q if it's not empty if not load_q.empty(): raise Exception(f"load_q not empty: {load_q.qsize()}", flush=True) + logger.info("Data Loading End") async def embed( @@ -200,7 +201,7 @@ async def embed( chan expects: (v_id, content, index_name) <- q.get() """ - logger.info("Reading from embed channel") + logger.info("Embedding Processing Start") async with asyncio.TaskGroup() as grp: # consume task queue while True: @@ -224,7 +225,7 @@ async def embed( except Exception: raise - logger.info(f"embed done") + logger.info("Embedding Processing End") async def extract( @@ -239,7 +240,7 @@ async def extract( chan expects: (chunk , chunk_id) <- q.get() """ - logger.info("Reading from extract channel") + logger.info("Entity Extration Start") # consume task queue async with asyncio.TaskGroup() as grp: while True: @@ -253,7 +254,7 @@ async def extract( except Exception: raise - logger.info(f"extract done") + logger.info("Entity Extration End") logger.info("closing upsert and embed chan") upsert_chan.close() @@ -268,7 +269,7 @@ async def stream_entities( """ Streams entity IDs from the grpah """ - logger.info("streaming entities") + logger.info("Entity Streaming Start") for i in range(ttl_batches): ids = await stream_ids(conn, "Entity", i, ttl_batches) if ids["error"]: @@ -280,7 +281,7 @@ async def stream_entities( if len(i) > 0: await entity_chan.put((i, "Entity")) - logger.info("stream_enities done") + logger.info("Entity Streaming End") # close the docs chan -- this function is the only sender logger.info("closing entities chan") entity_chan.close() @@ -298,6 +299,7 @@ async def resolve_entities( Copies edges between entities to their respective ResolvedEntities """ + logger.info("Entity Resolving Start") async with asyncio.TaskGroup() as grp: # for every entity while True: @@ -311,6 +313,7 @@ async def resolve_entities( break except Exception: raise + logger.info("Entity Resolving End") logger.info("closing upsert_chan") upsert_chan.close() logger.info("resolve_entities done") @@ -321,12 +324,12 @@ async def resolve_relationships( """ Copy RELATIONSHIP edges to RESOLVED_RELATIONSHIP """ - logger.info("Running ResolveRelationships") + logger.info("Relationship Resolving Start") async with tg_sem: res = await conn.runInstalledQuery( "ResolveRelationships" ) - logger.info("resolve_relationships done") + logger.info("Relationship Resolving End") async def communities(conn: AsyncTigerGraphConnection, comm_process_chan: Channel): """ diff --git a/graphrag-ui/src/components/CustomChatMessage.tsx b/graphrag-ui/src/components/CustomChatMessage.tsx index 33dea69..be577e7 100755 --- a/graphrag-ui/src/components/CustomChatMessage.tsx +++ b/graphrag-ui/src/components/CustomChatMessage.tsx @@ -111,7 +111,7 @@ export const CustomChatMessage: FC = ({ {showResult ? (
- Reasoning: + Reasoning:
{getReasoning(message)} = ({ )} ); -}; \ No newline at end of file +}; diff --git a/graphrag-ui/src/components/Interact.tsx b/graphrag-ui/src/components/Interact.tsx index a42793f..758fcef 100644 --- a/graphrag-ui/src/components/Interact.tsx +++ b/graphrag-ui/src/components/Interact.tsx @@ -110,4 +110,4 @@ export const Interactions: FC = ({ ) : null}
); -} \ No newline at end of file +} diff --git a/graphrag/app/agent.py b/graphrag/app/agent.py index 5fd0f5c..834eb48 100644 --- a/graphrag/app/agent.py +++ b/graphrag/app/agent.py @@ -54,12 +54,11 @@ def __init__( self.embedding_store = embedding_store self.mq2s = MapQuestionToSchema( - self.conn, self.llm.model, self.llm.map_question_schema_prompt + self.conn, self.llm ) self.gen_func = GenerateFunction( self.conn, - self.llm.model, - self.llm.generate_function_prompt, + self.llm, embedding_model, embedding_store, ) diff --git a/graphrag/app/agent/agent.py b/graphrag/app/agent/agent.py index b023ada..14f5df9 100644 --- a/graphrag/app/agent/agent.py +++ b/graphrag/app/agent/agent.py @@ -16,6 +16,7 @@ AWSBedrock, AzureOpenAI, GoogleVertexAI, + GoogleGenAI, Groq, HuggingFaceEndpoint, Ollama, @@ -63,14 +64,15 @@ def __init__( self.model_name = embedding_model.model_name self.embedding_model = embedding_model self.embedding_store = embedding_store + if self.embedding_store.conn.graphname != self.conn.graphname: + self.embedding_store.set_graphname(self.conn.graphname) self.mq2s = MapQuestionToSchema( - self.conn, self.llm.model, self.llm.map_question_schema_prompt + self.conn, self.llm ) self.gen_func = GenerateFunction( self.conn, - self.llm.model, - self.llm.generate_function_prompt, + self.llm, embedding_model, embedding_store, ) @@ -176,6 +178,9 @@ def make_agent(graphname, conn, use_cypher, ws: WebSocket = None, supportai_retr elif llm_config["completion_service"]["llm_service"].lower() == "vertexai": llm_service_name = "vertexai" llm_provider = GoogleVertexAI(llm_config["completion_service"]) + elif llm_config["completion_service"]["llm_service"].lower() == "genai": + llm_service_name = "genai" + llm_provider = GoogleGenAI(llm_config["completion_service"]) elif llm_config["completion_service"]["llm_service"].lower() == "bedrock": llm_service_name = "bedrock" llm_provider = AWSBedrock(llm_config["completion_service"]) diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py index f54be3d..c6d56e9 100644 --- a/graphrag/app/agent/agent_generation.py +++ b/graphrag/app/agent/agent_generation.py @@ -1,9 +1,11 @@ import logging from langchain.prompts import PromptTemplate from langchain_core.output_parsers import PydanticOutputParser +from langchain_community.callbacks.manager import get_openai_callback +from pydantic import BaseModel, Field from common.logs.logwriter import LogWriter from common.logs.log import req_id_cv -from langchain.pydantic_v1 import BaseModel, Field + logger = logging.getLogger(__name__) @@ -15,7 +17,7 @@ class TigerGraphAgentGenerator: def __init__(self, llm_model): self.llm = llm_model - def generate_answer(self, question: str, context: str) -> str: + def generate_answer(self, question: str, context: str) -> dict: """Generate an answer based on the question and context. Args: question: str: The question to generate an answer for. @@ -43,8 +45,16 @@ def generate_answer(self, question: str, context: str) -> str: # Chain rag_chain = prompt | self.llm.model | answer_parser - generation = rag_chain.invoke({"question": question, "context": context}) - + + usage_data = {} + with get_openai_callback() as cb: + generation = rag_chain.invoke({"question": question, "context": context}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + logger.info(f"generate_answer usage: {usage_data}") LogWriter.info(f"request_id={req_id_cv.get()} EXIT generate_answer") return generation diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py index 101a3c8..ae0d8e5 100644 --- a/graphrag/app/agent/agent_graph.py +++ b/graphrag/app/agent/agent_graph.py @@ -130,7 +130,7 @@ def map_question_to_schema(self, state): return state except MapQuestionToSchemaException as e: state["context"] = {"error": True} - if state["error_history"] is None: + if "error_history" not in state: state["error_history"] = [] state["error_history"].append({"error_message": str(e), "error_step": "generate_function"}) @@ -152,7 +152,7 @@ def generate_function(self, state): state["context"] = step except Exception as e: state["context"] = {"error": True} - if state["error_history"] is None: + if "error_history" not in state: state["error_history"] = [] state["error_history"].append({"error_message": str(e), "error_step": "generate_function"}) state["lookup_source"] = "inquiryai" @@ -311,7 +311,6 @@ def supportai_search(self, state): """ Run the agent supportai search. """ - if self.supportai_retriever == "hybridsearch": return self.hybrid_search(state) elif self.supportai_retriever == "similaritysearch": @@ -366,7 +365,7 @@ def generate_answer(self, state): citations = [re.sub(r"_chunk_\d+", "", x) for x in answer.citation] state["context"]["reasoning"] = list(set(citations)) - + try: resp = GraphRAGResponse( natural_language_response=answer.generated_answer, diff --git a/graphrag/app/agent/agent_hallucination_check.py b/graphrag/app/agent/agent_hallucination_check.py index 07533c4..6951c4a 100644 --- a/graphrag/app/agent/agent_hallucination_check.py +++ b/graphrag/app/agent/agent_hallucination_check.py @@ -1,9 +1,11 @@ import logging from langchain.prompts import PromptTemplate from langchain_core.output_parsers import PydanticOutputParser +from langchain_community.callbacks.manager import get_openai_callback + +from pydantic import BaseModel, Field from common.logs.logwriter import LogWriter from common.logs.log import req_id_cv -from langchain.pydantic_v1 import BaseModel, Field logger = logging.getLogger(__name__) @@ -46,6 +48,14 @@ def check_hallucination(self, generation: str, context: str) -> dict: # Chain rag_chain = prompt | self.llm.model | hallucination_parser - prediction = rag_chain.invoke({"context": context, "generation": generation}) + usage_data = {} + with get_openai_callback() as cb: + prediction = rag_chain.invoke({"context": context, "generation": generation}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + logger.info(f"check_hallucination usage: {usage_data}") LogWriter.info(f"request_id={req_id_cv.get()} EXIT check_hallucination") - return prediction \ No newline at end of file + return prediction diff --git a/graphrag/app/agent/agent_rewrite.py b/graphrag/app/agent/agent_rewrite.py index f5face6..ae83e54 100644 --- a/graphrag/app/agent/agent_rewrite.py +++ b/graphrag/app/agent/agent_rewrite.py @@ -2,9 +2,11 @@ import logging from langchain.prompts import PromptTemplate from langchain_core.output_parsers import PydanticOutputParser +from langchain_community.callbacks.manager import get_openai_callback + +from pydantic import BaseModel, Field from common.logs.log import req_id_cv from common.logs.logwriter import LogWriter -from langchain.pydantic_v1 import BaseModel, Field logger = logging.getLogger(__name__) @@ -29,7 +31,7 @@ def rewrite_question(self, question: str) -> str: re_write_prompt = PromptTemplate( template="""You are a question re-writer that converts an input question to a better version that is optimized \ -for AI agent question answering. Look at the initial and formulate an improved question. \n +for AI agent question answering. Look at the initial and formulate an improved question but avoid to add unnecessary context for entities. \n Here is the initial question: {question} Format your response in the following manner {format_instructions}""", @@ -43,6 +45,14 @@ def rewrite_question(self, question: str) -> str: # Chain question_rewriter = re_write_prompt | self.llm.model | rewrite_parser - generation = question_rewriter.invoke({"question": question}) + usage_data = {} + with get_openai_callback() as cb: + generation = question_rewriter.invoke({"question": question}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + logger.info(f"rewrite_question usage: {usage_data}") LogWriter.info(f"request_id={req_id_cv.get()} EXIT rewrite_question") return generation.rewritten_question diff --git a/graphrag/app/agent/agent_router.py b/graphrag/app/agent/agent_router.py index d969cb6..c86337c 100644 --- a/graphrag/app/agent/agent_router.py +++ b/graphrag/app/agent/agent_router.py @@ -1,11 +1,12 @@ from langchain.prompts import PromptTemplate from langchain_core.output_parsers import PydanticOutputParser +from langchain_community.callbacks.manager import get_openai_callback + +from pydantic import BaseModel, Field from common.logs.logwriter import LogWriter +from common.logs.log import req_id_cv from pyTigerGraph.pyTigerGraph import TigerGraphConnection import logging -from common.logs.log import req_id_cv - -from langchain.pydantic_v1 import BaseModel, Field logger = logging.getLogger(__name__) @@ -41,6 +42,14 @@ def route_question(self, question: str) -> str: ) question_router = prompt | self.llm.model | router_parser - res = question_router.invoke({"question": question, "v_types": v_types, "e_types": e_types}) + usage_data = {} + with get_openai_callback() as cb: + res = question_router.invoke({"question": question, "v_types": v_types, "e_types": e_types}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + logger.info(f"route_question usage: {usage_data}") LogWriter.info(f"request_id={req_id_cv.get()} EXIT route_question with {res}") return res diff --git a/graphrag/app/agent/agent_usefulness_check.py b/graphrag/app/agent/agent_usefulness_check.py index 226c7a9..425fc84 100644 --- a/graphrag/app/agent/agent_usefulness_check.py +++ b/graphrag/app/agent/agent_usefulness_check.py @@ -1,9 +1,11 @@ from langchain.prompts import PromptTemplate from langchain_core.output_parsers import PydanticOutputParser +from langchain_community.callbacks.manager import get_openai_callback + +from pydantic import BaseModel, Field from common.logs.logwriter import LogWriter -import logging from common.logs.log import req_id_cv -from langchain.pydantic_v1 import BaseModel, Field +import logging logger = logging.getLogger(__name__) @@ -46,6 +48,14 @@ def check_usefulness(self, question: str, answer: str) -> dict: # Chain rag_chain = prompt | self.llm.model | usefulness_parser - prediction = rag_chain.invoke({"generation": answer, "question": question}) + usage_data = {} + with get_openai_callback() as cb: + prediction = rag_chain.invoke({"generation": answer, "question": question}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + logger.info(f"check_usefulness usage: {usage_data}") LogWriter.info(f"request_id={req_id_cv.get()} EXIT check_usefulness") - return prediction \ No newline at end of file + return prediction diff --git a/graphrag/app/main.py b/graphrag/app/main.py index c75fc8c..9d34afb 100644 --- a/graphrag/app/main.py +++ b/graphrag/app/main.py @@ -44,6 +44,8 @@ logger = logging.getLogger(__name__) +logger.info("In main.py") + async def get_basic_auth_credentials(request: Request): auth_header = request.headers.get("Authorization") diff --git a/graphrag/app/routers/queryai.py b/graphrag/app/routers/queryai.py index edc4b25..f5f2794 100644 --- a/graphrag/app/routers/queryai.py +++ b/graphrag/app/routers/queryai.py @@ -56,7 +56,7 @@ def generate_cypher( natural_language_response=generated, answered_question=True, response_type="queryai", - query_sources={"graphname": graphname, "query": query.query} + query_sources={"graphname": graphname, "query": query.query}, ) return resp diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py index 4231af3..b664305 100644 --- a/graphrag/app/routers/ui.py +++ b/graphrag/app/routers/ui.py @@ -312,7 +312,6 @@ async def graph_query( prev_id = None while True: data = q - logger.info(f"Retrieving answer for chat \"{data}\"") # make message from data message = Message( @@ -397,7 +396,6 @@ async def chat( try: while True: data = await websocket.receive_text() - logger.info(f"Retrieving answer for chat \"{data}\"") # make message from data message = Message( diff --git a/graphrag/app/supportai/retrievers/BaseRetriever.py b/graphrag/app/supportai/retrievers/BaseRetriever.py index e85113c..647fb70 100644 --- a/graphrag/app/supportai/retrievers/BaseRetriever.py +++ b/graphrag/app/supportai/retrievers/BaseRetriever.py @@ -6,8 +6,9 @@ from langchain_core.output_parsers import StrOutputParser, PydanticOutputParser from langchain_core.prompts import ChatPromptTemplate, PromptTemplate -from langchain_core.pydantic_v1 import BaseModel, Field, validator from langchain.output_parsers import OutputFixingParser +from langchain_community.callbacks.manager import get_openai_callback + import re import logging @@ -67,7 +68,15 @@ def _question_to_keywords(self, question, top_k, verbose): chain = keyword_prompt | model | keyword_parser - answer = chain.invoke({"question": question}) + usage_data = {} + with get_openai_callback() as cb: + answer = chain.invoke({"question": question}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + self.logger.info(f"question_to_keywords usage: {usage_data}") if verbose: self.logger.info(f"Extracted keywords \"{answer}\" from question \"{question}\" by LLM") @@ -94,7 +103,15 @@ def _expand_question(self, question, top_k, verbose): chain = QUESTION_PROMPT | model | question_parser #chain = QUESTION_PROMPT | model - answer = chain.invoke({"question": question}) + usage_data = {} + with get_openai_callback() as cb: + answer = chain.invoke({"question": question}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + self.logger.info(f"expand_question usage: {usage_data}") if verbose: self.logger.info(f"Expanded question \"{question}\" from LLM: {answer}") @@ -119,7 +136,15 @@ def _generate_response(self, question, retrieved, verbose): chain = prompt | model | output_parser - generated = chain.invoke({"question": question, "sources": retrieved}) + usage_data = {} + with get_openai_callback() as cb: + generated = chain.invoke({"question": question, "sources": retrieved}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + self.logger.info(f"generate_response usage: {usage_data}") return {"response": generated, "retrieved": retrieved} @@ -144,7 +169,15 @@ def _hyde_embedding(self, text, str_mode: bool = False) -> str: chain = prompt | model | output_parser - generated = chain.invoke({"question": text}) + usage_data = {} + with get_openai_callback() as cb: + generated = chain.invoke({"question": text}) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + self.logger.info(f"hyde_embedding usage: {usage_data}") return self._generate_embedding(generated, str_mode) diff --git a/graphrag/app/supportai/retrievers/GraphRAGRetriever.py b/graphrag/app/supportai/retrievers/GraphRAGRetriever.py index 15a912e..e848c37 100644 --- a/graphrag/app/supportai/retrievers/GraphRAGRetriever.py +++ b/graphrag/app/supportai/retrievers/GraphRAGRetriever.py @@ -4,7 +4,8 @@ from langchain_core.output_parsers import PydanticOutputParser from langchain_core.prompts import PromptTemplate -from langchain_core.pydantic_v1 import BaseModel, Field, validator +from langchain_community.callbacks.manager import get_openai_callback + from supportai.retrievers import BaseRetriever from common.metrics.tg_proxy import TigerGraphConnectionProxy @@ -98,12 +99,19 @@ async def _generate_candidate(self, question, context): chain = ANSWER_PROMPT | model | answer_parser - answer = await chain.ainvoke( - { - "question": question, - "context": context, - } - ) + usage_data = {} + with get_openai_callback() as cb: + answer = await chain.ainvoke( + { + "question": question, + "context": context, + } + ) + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + self.logger.info(f"generate_candidate usage: {usage_data}") return answer diff --git a/graphrag/app/supportai/supportai.py b/graphrag/app/supportai/supportai.py index 160f604..fd963a0 100644 --- a/graphrag/app/supportai/supportai.py +++ b/graphrag/app/supportai/supportai.py @@ -5,6 +5,7 @@ from pyTigerGraph import TigerGraphConnection +from common.config import embedding_dimension from common.py_schemas.schemas import ( # GraphRAGResponse, CreateIngestConfig, @@ -53,6 +54,11 @@ def init_supportai(conn: TigerGraphConnection, graphname: str) -> tuple[dict, di file_path = "common/gsql/supportai/SupportAI_Schema_Native_Vector.gsql" with open(file_path, "r") as f: schema = f.read() + if embedding_dimension != 1536: + schema = schema.replace( + "dimension=1536", + f"dimension={embedding_dimension}", + ) schema_res += " " schema_res += conn.gsql( """USE GRAPH {}\n{}\nRUN SCHEMA_CHANGE JOB add_supportai_vector""".format( diff --git a/graphrag/app/tools/generate_cypher.py b/graphrag/app/tools/generate_cypher.py index 460047b..009f746 100644 --- a/graphrag/app/tools/generate_cypher.py +++ b/graphrag/app/tools/generate_cypher.py @@ -1,4 +1,5 @@ import logging +from langchain_community.callbacks.manager import get_openai_callback from langchain_core.output_parsers import StrOutputParser from langchain.prompts import PromptTemplate from langchain.tools import BaseTool @@ -12,8 +13,8 @@ class GenerateCypher(BaseTool): """GenerateCypher Tool. Tool to generate and execute the appropriate Cypher query for the question. """ - name = "GenerateCypher" - description = "Generates a Cypher query for the question." + name: str = "GenerateCypher" + description: str = "Generates a Cypher query for the question." conn: TigerGraphConnectionProxy = None llm: LLM = None @@ -93,7 +94,16 @@ def generate_cypher(self, question: str) -> str: logger.debug_pii("Prompt to LLM:\n" + PROMPT.invoke({"question": question, "schema": schema}).to_string()) chain = PROMPT | self.llm.model | StrOutputParser() - out = chain.invoke({"question": question, "schema": schema}).strip("```cypher").strip("```") + usage_data = {} + with get_openai_callback() as cb: + out = chain.invoke({"question": question, "schema": schema}).strip("```cypher").strip("```") + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + logger.info(f"generate_cypher usage: {usage_data}") + query_header = "USE GRAPH " + self.conn.graphname + " "+ "\n" + "INTERPRET OPENCYPHER QUERY () {" + "\n" query_footer = "\n}" return query_header + out + query_footer diff --git a/graphrag/app/tools/generate_function.py b/graphrag/app/tools/generate_function.py index 90040b5..77ecc12 100644 --- a/graphrag/app/tools/generate_function.py +++ b/graphrag/app/tools/generate_function.py @@ -2,13 +2,12 @@ import logging from typing import Dict, List, Optional, Type, Union -from langchain.chains import LLMChain from langchain.llms.base import LLM -from langchain.output_parsers import PydanticOutputParser +from langchain_core.output_parsers import PydanticOutputParser from langchain.prompts import PromptTemplate -from langchain.pydantic_v1 import BaseModel, Field, validator from langchain.tools import BaseTool from langchain.tools.base import ToolException +from langchain_community.callbacks.manager import get_openai_callback from common.embeddings.base_embedding_store import EmbeddingStore from common.embeddings.embedding_services import EmbeddingModel @@ -33,17 +32,16 @@ class GenerateFunction(BaseTool): Tool to generate and execute the appropriate function call for the question. """ - name = "GenerateFunction" - description = "Generates and executes a function call on the database. Always use MapQuestionToSchema before this tool." + name: str = "GenerateFunction" + description: str = "Generates and executes a function call on the database. Always use MapQuestionToSchema before this tool." conn: TigerGraphConnectionProxy = None llm: LLM = None - prompt: str = None handle_tool_error: bool = True embedding_model: EmbeddingModel = None embedding_store: EmbeddingStore = None args_schema: Type[MapQuestionToSchemaResponse] = MapQuestionToSchemaResponse - def __init__(self, conn, llm, prompt, embedding_model, embedding_store): + def __init__(self, conn, llm, embedding_model, embedding_store): """Initialize GenerateFunction. Args: conn (TigerGraphConnection): @@ -61,7 +59,6 @@ def __init__(self, conn, llm, prompt, embedding_model, embedding_store): logger.debug(f"request_id={req_id_cv.get()} GenerateFunction instantiated") self.conn = conn self.llm = llm - self.prompt = prompt self.embedding_model = embedding_model self.embedding_store = embedding_store @@ -122,7 +119,7 @@ def _run( func_parser = PydanticOutputParser(pydantic_object=GenerateFunctionResponse) PROMPT = PromptTemplate( - template=self.prompt, + template=self.llm.generate_function_prompt, input_variables=[ "question", "vertex_types", @@ -170,31 +167,37 @@ def _run( LogWriter.warning(f"request_id={req_id_cv.get()} WARN no documents found") raise NoDocumentsFoundException - inputs = [ - { - "question": question, - "vertex_types": target_vertex_types, - "edge_types": target_edge_types, - "vertex_attributes": target_vertex_attributes, - "vertex_ids": target_vertex_ids, - "edge_attributes": target_edge_attributes, - "doc1": docs[0].page_content, - "doc2": docs[1].page_content if len(docs) > 1 else "", - "doc3": docs[2].page_content if len(docs) > 2 else "", - "doc4": docs[3].page_content if len(docs) > 3 else "", - "doc5": docs[4].page_content if len(docs) > 4 else "", - "doc6": docs[5].page_content if len(docs) > 5 else "", - "doc7": docs[6].page_content if len(docs) > 6 else "", - "doc8": docs[7].page_content if len(docs) > 7 else "", - } - ] + inputs = { + "question": question, + "vertex_types": target_vertex_types, + "edge_types": target_edge_types, + "vertex_attributes": target_vertex_attributes, + "vertex_ids": target_vertex_ids, + "edge_attributes": target_edge_attributes, + "doc1": docs[0].page_content, + "doc2": docs[1].page_content if len(docs) > 1 else "", + "doc3": docs[2].page_content if len(docs) > 2 else "", + "doc4": docs[3].page_content if len(docs) > 3 else "", + "doc5": docs[4].page_content if len(docs) > 4 else "", + "doc6": docs[5].page_content if len(docs) > 5 else "", + "doc7": docs[6].page_content if len(docs) > 6 else "", + "doc8": docs[7].page_content if len(docs) > 7 else "", + } logger.debug(f"request_id={req_id_cv.get()} retrieved documents={docs}") - chain = LLMChain(llm=self.llm, prompt=PROMPT) - generated = chain.apply(inputs)[0]["text"] + chain = PROMPT | self.llm.model | func_parser + usage_data = {} + with get_openai_callback() as cb: + generated = chain.invoke(**inputs) + + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + logger.info(f"generate_function usage: {usage_data}") + logger.debug(f"request_id={req_id_cv.get()} generated function") - generated = func_parser.invoke(generated) try: parsed_func = validate_function_call( self.conn, generated.connection_func_call, valid_function_calls diff --git a/graphrag/app/tools/map_question_to_schema.py b/graphrag/app/tools/map_question_to_schema.py index 80013f7..13c62a7 100644 --- a/graphrag/app/tools/map_question_to_schema.py +++ b/graphrag/app/tools/map_question_to_schema.py @@ -1,10 +1,10 @@ from langchain.tools import BaseTool -from langchain.llms.base import LLM from langchain.tools.base import ToolException -from langchain.chains import LLMChain +from langchain.llms.base import LLM from langchain.prompts import PromptTemplate -from langchain.output_parsers import PydanticOutputParser -from langchain.pydantic_v1 import BaseModel, Field, validator +from langchain_core.output_parsers import PydanticOutputParser +from langchain_community.callbacks.manager import get_openai_callback + from common.metrics.tg_proxy import TigerGraphConnectionProxy from common.py_schemas import MapQuestionToSchemaResponse, MapAttributeToAttributeResponse from typing import List, Dict @@ -22,14 +22,14 @@ class MapQuestionToSchema(BaseTool): Tool to map questions to their datatypes in the database. Should be executed before GenerateFunction. """ - name = "MapQuestionToSchema" - description = "Always run first to map the query to the graph's schema. GenerateFunction before using MapQuestionToSchema" - conn: "TigerGraphConnectionProxy" = None + name: str = "MapQuestionToSchema" + description: str = "Always run first to map the query to the graph's schema. GenerateFunction before using MapQuestionToSchema" + conn: TigerGraphConnectionProxy = None llm: LLM = None prompt: str = None handle_tool_error: bool = True - def __init__(self, conn, llm, prompt): + def __init__(self, conn, llm): """Initialize MapQuestionToSchema. Args: conn (TigerGraphConnectionProxy): @@ -43,7 +43,7 @@ def __init__(self, conn, llm, prompt): logger.debug(f"request_id={req_id_cv.get()} MapQuestionToSchema instantiated") self.conn = conn self.llm = llm - self.prompt = prompt + def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: """Run the tool. @@ -55,7 +55,7 @@ def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: parser = PydanticOutputParser(pydantic_object=MapQuestionToSchemaResponse) RESTATE_QUESTION_PROMPT = PromptTemplate( - template=self.prompt, + template=self.llm.map_question_schema_prompt, input_variables=[ "question", "conversation", @@ -67,7 +67,7 @@ def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: partial_variables={"format_instructions": parser.get_format_instructions()}, ) - restate_chain = LLMChain(llm=self.llm, prompt=RESTATE_QUESTION_PROMPT) + restate_chain = RESTATE_QUESTION_PROMPT | self.llm.model | parser vertices = self.conn.getVertexTypes() edges = self.conn.getEdgeTypes() @@ -86,23 +86,23 @@ def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: edge_info = {"edge": edge, "source": source_vertex, "target": target_vertex} edges_info.append(edge_info) - restate_q = restate_chain.apply( - [ + usage_data = {} + with get_openai_callback() as cb: + parsed_q = restate_chain.invoke( { "vertices": vertices, "verticesAttrs": vertices_info, "edges": edges, "edgesInfo": edges_info, "question": query, - "conversation": conversation + "conversation": conversation, } - ] - )[0]["text"] - - logger.debug(f"request_id={req_id_cv.get()} MapQuestionToSchema applied") - # logger.info(f"restate_q: {restate_q}") - - parsed_q = parser.invoke(restate_q) + ) + usage_data["input_tokens"] = cb.prompt_tokens + usage_data["output_tokens"] = cb.completion_tokens + usage_data["total_tokens"] = cb.total_tokens + usage_data["cost"] = cb.total_cost + # logger.info(f"parsed_q: {parsed_q}") logger.debug_pii( f"request_id={req_id_cv.get()} MapQuestionToSchema parsed for question={query} into normalized_form={parsed_q}" @@ -124,38 +124,46 @@ def _run(self, query: str, conversation: List[Dict[str, str]]) -> str: }, ) - attr_map_chain = LLMChain(llm=self.llm, prompt=ATTR_MAP_PROMPT) - for vertex in parsed_q.target_vertex_attributes.keys(): - map_attr = attr_map_chain.apply( - [ - { - "parsed_attrs": parsed_q.target_vertex_attributes[vertex], - "real_attrs": [attr[0] for attr in self.conn.getVertexAttrs(vertex)], - } + attr_map_chain = ATTR_MAP_PROMPT | self.llm.model | attr_parser + if parsed_q.target_vertex_attributes: + for vertex in parsed_q.target_vertex_attributes.keys(): + with get_openai_callback() as cb: + parsed_map = attr_map_chain.invoke( + { + "parsed_attrs": parsed_q.target_vertex_attributes[vertex], + "real_attrs": [attr[0] for attr in self.conn.getVertexAttrs(vertex)], + } + ).attr_map + usage_data["input_tokens"] += cb.prompt_tokens + usage_data["output_tokens"] += cb.completion_tokens + usage_data["total_tokens"] += cb.total_tokens + usage_data["cost"] += cb.total_cost + parsed_q.target_vertex_attributes[vertex] = [ + parsed_map.get(x) for x in list(parsed_q.target_vertex_attributes[vertex]) ] - )[0]["text"] - parsed_map = attr_parser.invoke(map_attr).attr_map - parsed_q.target_vertex_attributes[vertex] = [ - parsed_map.get(x) for x in list(parsed_q.target_vertex_attributes[vertex]) - ] - - logger.debug(f"request_id={req_id_cv.get()} MapVertexAttributes applied") - - for edge in parsed_q.target_edge_attributes.keys(): - map_attr = attr_map_chain.apply( - [ - { - "parsed_attrs": parsed_q.target_edge_attributes[edge], - "real_attrs": self.conn.getEdgeAttrs(edge), - } + + logger.debug(f"request_id={req_id_cv.get()} MapVertexAttributes applied") + + if parsed_q.target_edge_attributes: + for edge in parsed_q.target_edge_attributes.keys(): + with get_openai_callback() as cb: + parsed_map = attr_map_chain.invoke( + { + "parsed_attrs": parsed_q.target_edge_attributes[edge], + "real_attrs": self.conn.getEdgeAttrs(edge), + } + ).attr_map + usage_data["input_tokens"] += cb.prompt_tokens + usage_data["output_tokens"] += cb.completion_tokens + usage_data["total_tokens"] += cb.total_tokens + usage_data["cost"] += cb.total_cost + parsed_q.target_edge_attributes[edge] = [ + parsed_map[x] for x in list(parsed_q.target_edge_attributes[edge]) ] - )[0]["text"] - parsed_map = attr_parser.invoke(map_attr).attr_map - parsed_q.target_edge_attributes[edge] = [ - parsed_map[x] for x in list(parsed_q.target_edge_attributes[edge]) - ] - logger.debug(f"request_id={req_id_cv.get()} MapEdgeAttributes applied") + logger.debug(f"request_id={req_id_cv.get()} MapEdgeAttributes applied") + + logger.info(f"map_question_to_schema usage: {usage_data}") try: validate_schema( diff --git a/graphrag/app/tools/validation_utils.py b/graphrag/app/tools/validation_utils.py index f03765a..6b18b35 100644 --- a/graphrag/app/tools/validation_utils.py +++ b/graphrag/app/tools/validation_utils.py @@ -25,47 +25,51 @@ class InvalidFunctionCallException(Exception): def validate_schema(conn, v_types, e_types, v_attrs, e_attrs): LogWriter.info(f"request_id={req_id_cv.get()} ENTRY validate_schema()") - vertices = conn.getVertexTypes() - edges = conn.getEdgeTypes() - for v in v_types: - logger.debug( - f"request_id={req_id_cv.get()} validate_schema() validating vertex_type={v}" - ) - if v in vertices: - attrs = [x["AttributeName"] for x in conn.getVertexType(v)["Attributes"]] - for attr in v_attrs.get(v, []): - if attr not in attrs and attr != "": - if attr is None: - attr = "None" - raise MapQuestionToSchemaException( - f"{attr} is not found for {v} in the data schema. Run MapQuestionToSchema to validate schema." - ) - else: - if v is None: - v = "None" - raise MapQuestionToSchemaException( - f"{v} is not found in the data schema. Run MapQuestionToSchema to validate schema." + if v_types: + vertices = conn.getVertexTypes() + for v in v_types: + logger.debug( + f"request_id={req_id_cv.get()} validate_schema() validating vertex_type={v}" ) + if v in vertices: + if v_attrs: + attrs = [x["AttributeName"] for x in conn.getVertexType(v)["Attributes"]] + for attr in v_attrs.get(v, []): + if attr not in attrs and attr != "": + if attr is None: + attr = "None" + raise MapQuestionToSchemaException( + f"{attr} is not found for {v} in the data schema. Run MapQuestionToSchema to validate schema." + ) + else: + if v is None: + v = "None" + raise MapQuestionToSchemaException( + f"{v} is not found in the data schema. Run MapQuestionToSchema to validate schema." + ) - for e in e_types: - logger.debug( - f"request_id={req_id_cv.get()} validate_schema() validating edge_type={e}" - ) - if e in edges: - attrs = [x["AttributeName"] for x in conn.getEdgeType(e)["Attributes"]] - for attr in e_attrs.get(e, []): - if attr not in attrs and attr != "": - if attr is None: - attr = "None" - raise MapQuestionToSchemaException( - f"{attr} is not found for {e} in the data schema. Run MapQuestionToSchema to validate schema." - ) - else: - if e is None: - e = "None" - raise MapQuestionToSchemaException( - f"{e} is not found in the data schema. Run MapQuestionToSchema to validate schema." + if e_types: + edges = conn.getEdgeTypes() + for e in e_types: + logger.debug( + f"request_id={req_id_cv.get()} validate_schema() validating edge_type={e}" ) + if e in edges: + if e_attrs: + attrs = [x["AttributeName"] for x in conn.getEdgeType(e)["Attributes"]] + for attr in e_attrs.get(e, []): + if attr not in attrs and attr != "": + if attr is None: + attr = "None" + raise MapQuestionToSchemaException( + f"{attr} is not found for {e} in the data schema. Run MapQuestionToSchema to validate schema." + ) + else: + if e is None: + e = "None" + raise MapQuestionToSchemaException( + f"{e} is not found in the data schema. Run MapQuestionToSchema to validate schema." + ) LogWriter.info(f"request_id={req_id_cv.get()} EXIT validate_schema()") return True diff --git a/report-service/app/report_agent/agent.py b/report-service/app/report_agent/agent.py index 78d89d5..b04ce9f 100644 --- a/report-service/app/report_agent/agent.py +++ b/report-service/app/report_agent/agent.py @@ -2,10 +2,10 @@ from common.llm_services.base_llm import LLM_Model from common.py_schemas.tool_io_schemas import ReportSections, ReportSection - from langchain.prompts import PromptTemplate from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser -from langchain.pydantic_v1 import BaseModel, Field +from pydantic import BaseModel, Field + class Report(BaseModel): report: str = Field(description="A drareport based on the answers received from the data analyst. Include citations by adding `[x]` where x is the function call that was used to determine the answer.")