-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* Dependencies updated * CORs middleware added * Neo4j exception middleware added * Replaced deprecated LLMChain implementation * Vector chain simplified to use RetrievalQA chain
- Loading branch information
Showing
7 changed files
with
1,244 additions
and
586 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,60 +1,113 @@ | ||
| from __future__ import annotations | ||
| from typing import Union | ||
| from app.graph_chain import graph_chain, CYPHER_GENERATION_PROMPT | ||
| from app.vector_chain import vector_chain, VECTOR_PROMPT | ||
| from app.simple_agent import simple_agent_chain | ||
| from fastapi import FastAPI | ||
| from typing import Union, Optional | ||
| from fastapi import FastAPI, Request, Response | ||
| from fastapi.middleware.cors import CORSMiddleware | ||
| from starlette.middleware.base import BaseHTTPMiddleware | ||
| from pydantic import BaseModel, Field | ||
| from neo4j import exceptions | ||
| import logging | ||
|
|
||
|
|
||
| class ApiChatPostRequest(BaseModel): | ||
| message: str = Field(..., description='The chat message to send') | ||
| mode: str = Field('agent', description='The mode of the chat message. Current options are: "vector", "graph", "agent". Default is "agent"') | ||
| message: str = Field(..., description="The chat message to send") | ||
| mode: str = Field( | ||
| "agent", | ||
| description='The mode of the chat message. Current options are: "vector", "graph", "agent". Default is "agent"', | ||
| ) | ||
|
|
||
|
|
||
| class ApiChatPostResponse(BaseModel): | ||
| message: Optional[str] = Field(None, description='The chat message response') | ||
| response: str | ||
|
|
||
|
|
||
| class Neo4jExceptionMiddleware(BaseHTTPMiddleware): | ||
| async def dispatch(self, request: Request, call_next): | ||
| try: | ||
| response = await call_next(request) | ||
| return response | ||
| except exceptions.AuthError as e: | ||
| msg = f"Neo4j Authentication Error: {e}" | ||
| logging.warning(msg) | ||
| return Response(content=msg, status_code=400, media_type="text/plain") | ||
| except exceptions.ServiceUnavailable as e: | ||
| msg = f"Neo4j Database Unavailable Error: {e}" | ||
| logging.warning(msg) | ||
| return Response(content=msg, status_code=400, media_type="text/plain") | ||
| except Exception as e: | ||
| msg = f"Neo4j Uncaught Exception: {e}" | ||
| logging.error(msg) | ||
| return Response(content=msg, status_code=400, media_type="text/plain") | ||
|
|
||
|
|
||
| # Allowed CORS origins | ||
| origins = [ | ||
| "http://127.0.0.1:8000", # Alternative localhost address | ||
| "http://localhost:8000", | ||
| ] | ||
|
|
||
| app = FastAPI() | ||
|
|
||
| # Add CORS middleware to allow cross-origin requests | ||
| app.add_middleware( | ||
| CORSMiddleware, | ||
| allow_origins=origins, | ||
| allow_credentials=True, | ||
| allow_methods=["*"], | ||
| allow_headers=["*"], | ||
| ) | ||
| # Add Neo4j exception handling middleware | ||
| app.add_middleware(Neo4jExceptionMiddleware) | ||
|
|
||
| app = FastAPI() | ||
|
|
||
| @app.post( | ||
| '/api/chat', | ||
| "/api/chat", | ||
| response_model=None, | ||
| responses={'201': {'model': ApiChatPostResponse}}, | ||
| tags=['chat'], | ||
| responses={"201": {"model": ApiChatPostResponse}}, | ||
| tags=["chat"], | ||
| ) | ||
| def send_chat_message(body: ApiChatPostRequest) -> Union[None, ApiChatPostResponse]: | ||
| async def send_chat_message(body: ApiChatPostRequest): | ||
| """ | ||
| Send a chat message | ||
| """ | ||
|
|
||
| question = body.message | ||
|
|
||
| v_response = vector_chain().invoke( | ||
| {"question":question}, | ||
| prompt = VECTOR_PROMPT, | ||
| return_only_outputs = True | ||
| ) | ||
| g_response = graph_chain().invoke( | ||
| {"query":question}, | ||
| prompt = CYPHER_GENERATION_PROMPT, | ||
| return_only_outputs = True | ||
| ) | ||
|
|
||
| if body.mode == 'vector': | ||
| # Simple exception check. See https://neo4j.com/docs/api/python-driver/current/api.html#errors for full set of driver exceptions | ||
|
|
||
| if body.mode == "vector": | ||
| # Return only the Vector answer | ||
| v_response = vector_chain().invoke( | ||
| {"query": question}, prompt=VECTOR_PROMPT, return_only_outputs=True | ||
| ) | ||
| response = v_response | ||
| elif body.mode == 'graph': | ||
| elif body.mode == "graph": | ||
| # Return only the Graph (text2Cypher) answer | ||
| response = g_response | ||
| g_response = graph_chain().invoke( | ||
| {"query": question}, | ||
| prompt=CYPHER_GENERATION_PROMPT, | ||
| return_only_outputs=True, | ||
| ) | ||
| response = g_response["result"] | ||
| else: | ||
| # Return an answer from a chain that composites both the Vector and Graph responses | ||
| response = simple_agent_chain().invoke({ | ||
| "question":question, | ||
| "vector_result":v_response, | ||
| "graph_result":g_response | ||
| })["text"] | ||
|
|
||
| return f"{response}", 200 | ||
| # Return both vector + graph answers | ||
| v_response = vector_chain().invoke( | ||
| {"query": question}, prompt=VECTOR_PROMPT, return_only_outputs=True | ||
| ) | ||
| g_response = graph_chain().invoke( | ||
| {"query": question}, | ||
| prompt=CYPHER_GENERATION_PROMPT, | ||
| return_only_outputs=True, | ||
| )["result"] | ||
|
|
||
| # Synthesize a composite of both the Vector and Graph responses | ||
| response = simple_agent_chain().invoke( | ||
| { | ||
| "question": question, | ||
| "vector_result": v_response, | ||
| "graph_result": g_response, | ||
| } | ||
| ) | ||
|
|
||
| return response, 200 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.