基于图的元数据过滤以提高RAG应用中的向量搜索性能

Neo4j Browser URI: https://demo.neo4jlabs.com:7473/browser/username: companiespassword: companiesdatabase: companies

import os
os.environ["OPENAI_API_KEY"] = "sk-"os.environ["NEO4J_URI"] = "neo4j+s://demo.neo4jlabs.com"os.environ["NEO4J_USERNAME"] = "companies"os.environ["NEO4J_PASSWORD"] = "companies"os.environ["NEO4J_DATABASE"] = "companies"
embeddings = OpenAIEmbeddings()graph = Neo4jGraph()vector_index = Neo4jVector.from_existing_index(    embeddings,     index_name="news")

def get_organization_news(    topic: Optional[str] = None,    organization: Optional[str] = None,    country: Optional[str] = None,    sentiment: Optional[str] = None,) -> str:    # If there is no prefiltering, we can use vector index    if topic and not organization and not country and not sentiment:        return vector_index.similarity_search(topic)    # Uses parallel runtime where available    base_query = (        "CYPHER runtime = parallel parallelRuntimeSupport=all "        "MATCH (c:Chunk)<-[:HAS_CHUNK]-(a:Article) WHERE "    )    where_queries = []    params = {"k": 5} # Define the number of text chunks to retrieve

if organization:    # Map to database    candidates = get_candidates(organization)    if len(candidates) > 1:  # Ask for follow up if too many options        return (         "Ask a follow up question which of the available organizations "         f"did the user mean. Available options: {candidates}"        )    where_queries.append(        "EXISTS {(a)-[:MENTIONS]->(:Organization {name: $organization})}"    )    params["organization"] = candidates[0]

if country:    # No need to disambiguate    where_queries.append(        "EXISTS {(a)-[:MENTIONS]->(:Organization)-[:IN_CITY]->()-[:IN_COUNTRY]->(:Country {name: $country})}"    )    params["country"] = country

if sentiment:    if sentiment == "positive":        where_queries.append("a.sentiment > $sentiment")        params["sentiment"] = 0.5    else:        where_queries.append("a.sentiment < $sentiment")        params["sentiment"] = -0.5

if topic:  # Do vector comparison    vector_snippet = (        " WITH c, a, vector.similarity.cosine(c.embedding,$embedding) AS score "        "ORDER BY score DESC LIMIT toInteger($k) "    )    params["embedding"] = embeddings.embed_query(topic)else:  # Just return the latest data    vector_snippet = " WITH c, a ORDER BY a.date DESC LIMIT toInteger($k) "

return_snippet = "RETURN '#title ' + a.title + '\n#date ' + toString(a.date) + '\n#text ' + c.text AS output"
complete_query = (    base_query + " AND ".join(where_queries) + vector_snippet + return_snippet)
# Retrieve information from the databasedata = graph.query(complete_query, params)print(f"Cypher: {complete_query}\n")# Safely remove embedding before printingparams.pop('embedding', None)print(f"Parameters: {params}")return "###Article: ".join([el["output"] for el in data])

get_organization_news(  organization='neo4j',  sentiment='positive',  topic='remote work')
# Cypher: CYPHER runtime = parallel parallelRuntimeSupport=all# MATCH (c:Chunk)<-[:HAS_CHUNK]-(a:Article) WHERE # EXISTS {(a)-[:MENTIONS]->(:Organization {name: $organization})} AND # a.sentiment > $sentiment # WITH c, a, vector.similarity.cosine(c.embedding,$embedding) AS score # ORDER BY score DESC LIMIT toInteger($k) # RETURN '#title ' + a.title + '\ndate ' + toString(a.date) + '\ntext ' + c.text AS output
# Parameters: {'k': 5, 'organization': 'Neo4j', 'sentiment': 0.5}

fewshot_examples = """{Input:What are the health benefits for Google employees in the news? Query: Health benefits}{Input: What is the latest positive news about Google? Query: None}{Input: Are there any news about VertexAI regarding Google? Query: VertexAI}{Input: Are there any news about new products regarding Google? Query: new products}"""
class NewsInput(BaseModel):    topic: Optional[str] = Field(        description="Any specific information or topic besides organization, country, and sentiment that the user is interested in. Here are some examples: "        + fewshot_examples    )    organization: Optional[str] = Field(        description="Organization that the user wants to find information about"    )    country: Optional[str] = Field(        description="Country of organizations that the user is interested in. Use full names like United States of America and France."    )    sentiment: Optional[str] = Field(        description="Sentiment of articles", enum=["positive", "negative"]    )

class NewsTool(BaseTool):    name = "NewsInformation"    description = (        "useful for when you need to find relevant information in the news"    )    args_schema: Type[BaseModel] = NewsInput
    def _run(        self,        topic: Optional[str] = None,        organization: Optional[str] = None,        country: Optional[str] = None,        sentiment: Optional[str] = None,        run_manager: Optional[CallbackManagerForToolRun] = None,    ) -> str:        """Use the tool."""        return get_organization_news(topic, organization, country, sentiment)

llm = ChatOpenAI(temperature=0, model="gpt-4-turbo", streaming=True)tools = [NewsTool()]
llm_with_tools = llm.bind(functions=[format_tool_to_openai_function(t) for t in tools])
prompt = ChatPromptTemplate.from_messages(    [        (            "system",            "You are a helpful assistant that finds information about movies "            " and recommends them. If tools require follow up questions, "            "make sure to ask the user for clarification. Make sure to include any "            "available options that need to be clarified in the follow up questions "            "Do only the things the user specifically requested. ",        ),        MessagesPlaceholder(variable_name="chat_history"),        ("user", "{input}"),        MessagesPlaceholder(variable_name="agent_scratchpad"),    ])
agent = (    {        "input": lambda x: x["input"],        "chat_history": lambda x: _format_chat_history(x["chat_history"])        if x.get("chat_history")        else [],        "agent_scratchpad": lambda x: format_to_openai_function_messages(            x["intermediate_steps"]        ),    }    | prompt    | llm_with_tools    | OpenAIFunctionsAgentOutputParser())
agent_executor = AgentExecutor(agent=agent, tools=tools)

agent_executor.invoke(  {"input": "What are some positive news regarding neo4j?"})
# Cypher: CYPHER runtime = parallel parallelRuntimeSupport=all # MATCH (c:Chunk)<-[:HAS_CHUNK]-(a:Article) WHERE # EXISTS {(a)-[:MENTIONS]->(:Organization {name: $organization})} AND # a.sentiment > $sentiment WITH c, a # ORDER BY a.date DESC LIMIT toInteger($k) # RETURN '#title ' + a.title + 'date ' + toString(a.date) + 'text ' + c.text AS output# Parameters: {'k': 5, 'organization': 'Neo4j', 'sentiment': 0.5}

agent_executor.invoke(   {"input": "What are some of the latest negative news about employee happiness for companies from France?"})
# Cypher: CYPHER runtime = parallel parallelRuntimeSupport=all # MATCH (c:Chunk)<-[:HAS_CHUNK]-(a:Article) WHERE # EXISTS {(a)-[:MENTIONS]->(:Organization)-[:IN_CITY]->()-[:IN_COUNTRY]->(:Country {name: $country})} AND # a.sentiment < $sentiment # WITH c, a, vector.similarity.cosine(c.embedding,$embedding) AS score # ORDER BY score DESC LIMIT toInteger($k) # RETURN '#title ' + a.title + 'date ' + toString(a.date) + 'text ' + c.text AS output# Parameters: {'k': 5, 'country': 'France', 'sentiment': -0.5, 'topic': 'employee happiness'}