refactor: standardize imports, update sample queries for example and index type payload in vector db

Shrijeeth-Suresh · Shrijeeth-Suresh · commit 9e956ded8f5f · 2025-07-05T19:45:45.000+05:30
diff --git a/backend/agents/query_agent/agent.py b/backend/agents/query_agent/agent.py
@@ -1,12 +1,13 @@
 from typing import AsyncGenerator
 
-from config import get_settings
 from google.adk.agents import BaseAgent, LoopAgent, SequentialAgent
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.events import Event, EventActions
 from google.adk.models.lite_llm import LiteLlm
 from google.adk.tools.agent_tool import AgentTool
 
+from config import get_settings
+
 from .sub_agents import (
     column_selection_agent,
     query_explanation_agent,
diff --git a/backend/agents/query_agent/sub_agents/column_selection_agent/agent.py b/backend/agents/query_agent/sub_agents/column_selection_agent/agent.py
@@ -1,8 +1,9 @@
-from config import get_settings
 from google.adk.agents import LlmAgent
 from google.adk.models.lite_llm import LiteLlm
 from google.genai import types
 
+from config import get_settings
+
 from .prompts import agent_description, agent_instruction
 
 AGENT_MODEL = LiteLlm(
diff --git a/backend/agents/query_agent/sub_agents/query_explanation_agent/agent.py b/backend/agents/query_agent/sub_agents/query_explanation_agent/agent.py
@@ -1,8 +1,9 @@
-from config import get_settings
 from google.adk.agents import LlmAgent
 from google.adk.models.lite_llm import LiteLlm
 from google.genai import types
 
+from config import get_settings
+
 from . import tools
 from .prompts import agent_description, agent_instruction
 
diff --git a/backend/agents/query_agent/sub_agents/query_explanation_agent/tools.py b/backend/agents/query_agent/sub_agents/query_explanation_agent/tools.py
@@ -1,11 +1,12 @@
 import json
 
-from config import get_settings
 from google.adk.tools import FunctionTool
 from litellm import embedding
 from qdrant_client import QdrantClient
 from qdrant_client.models import FieldCondition, Filter, MatchValue
 
+from config import get_settings
+
 
 def get_similar_queries(user_input: str, score_threshold: float = 0.3) -> str:
     """
diff --git a/backend/agents/query_agent/sub_agents/query_formation_agent/agent.py b/backend/agents/query_agent/sub_agents/query_formation_agent/agent.py
@@ -1,8 +1,9 @@
-from config import get_settings
 from google.adk.agents import LlmAgent
 from google.adk.models.lite_llm import LiteLlm
 from google.genai import types
 
+from config import get_settings
+
 from .prompts import agent_description, agent_instruction
 
 AGENT_MODEL = LiteLlm(
diff --git a/backend/agents/query_agent/sub_agents/query_validation_agent/agent.py b/backend/agents/query_agent/sub_agents/query_validation_agent/agent.py
@@ -1,8 +1,9 @@
-from config import get_settings
 from google.adk.agents import LlmAgent
 from google.adk.models.lite_llm import LiteLlm
 from google.genai import types
 
+from config import get_settings
+
 from . import tools
 from .prompts import agent_description, agent_instruction
 
diff --git a/backend/agents/query_agent/sub_agents/query_validation_agent/tools.py b/backend/agents/query_agent/sub_agents/query_validation_agent/tools.py
@@ -1,7 +1,8 @@
-from config import get_settings
 from google.adk.tools import FunctionTool
 from mysql.connector import Error as MySQLError
 from psycopg2 import Error as PostgresError
+
+from config import get_settings
 from utils.helpers import (
     get_validate_db_mysql_connection,
     get_validate_db_postgres_connection,
diff --git a/backend/agents/query_agent/sub_agents/table_agent/agent.py b/backend/agents/query_agent/sub_agents/table_agent/agent.py
@@ -1,8 +1,9 @@
-from config import get_settings
 from google.adk.agents import LlmAgent
 from google.adk.models.lite_llm import LiteLlm
 from google.genai import types
 
+from config import get_settings
+
 from . import tools
 from .prompts import agent_description, agent_instruction
 
diff --git a/backend/agents/query_agent/sub_agents/table_agent/tools.py b/backend/agents/query_agent/sub_agents/table_agent/tools.py
@@ -1,11 +1,12 @@
 import json
 
-from config import get_settings
 from google.adk.tools import FunctionTool
 from litellm import embedding
 from qdrant_client import QdrantClient
 from qdrant_client.models import FieldCondition, Filter, MatchValue
 
+from config import get_settings
+
 
 def get_similar_tables(
     user_input: str, similar_queries: str, score_threshold: float = 0.3
@@ -37,11 +38,15 @@ def get_similar_tables(
         collection_name=get_settings().QDRANT_COLLECTION_NAME,
         query=embed.data[0].embedding,
         query_filter=Filter(
-            must=[
+            should=[
                 FieldCondition(
                     key="type",
                     match=MatchValue(value="TABLE"),
-                )
+                ),
+                FieldCondition(
+                    key="type",
+                    match=MatchValue(value="VIEW"),
+                ),
             ]
         ),
         score_threshold=score_threshold,
diff --git a/backend/agents/sql_data_generation_agent/sub_agents/sql_query_generation_agent/agent.py b/backend/agents/sql_data_generation_agent/sub_agents/sql_query_generation_agent/agent.py
@@ -1,8 +1,9 @@
-from config import get_settings
 from google.adk.agents import LlmAgent
 from google.adk.models.lite_llm import LiteLlm
 from google.genai import types
 
+from config import get_settings
+
 from .prompts import agent_description, agent_instruction
 from .types import SampleQuerySchema
 
diff --git a/backend/agents/sql_table_describer_agent/sub_agents/table_description_agent/agent.py b/backend/agents/sql_table_describer_agent/sub_agents/table_description_agent/agent.py
@@ -1,8 +1,9 @@
-from config import get_settings
 from google.adk.agents import LlmAgent
 from google.adk.models.lite_llm import LiteLlm
 from google.genai import types
 
+from config import get_settings
+
 from .prompts import agent_description, agent_instruction
 from .types import TableDescriptionSchema
 
diff --git a/backend/alembic/env.py b/backend/alembic/env.py
@@ -2,10 +2,11 @@
 import sys
 from logging.config import fileConfig
 
+from sqlalchemy import engine_from_config, pool
+
 from alembic import context
 from config import get_settings
 from models import Base
-from sqlalchemy import engine_from_config, pool
 
 # import Base and User from models package (Base is now unified)
 
diff --git a/backend/alembic/versions/5c3984fbf0a1_initial_schema.py b/backend/alembic/versions/5c3984fbf0a1_initial_schema.py
@@ -9,6 +9,7 @@
 from typing import Sequence, Union
 
 import sqlalchemy as sa
+
 from alembic import op
 
 # revision identifiers, used by Alembic.
diff --git a/backend/auth.py b/backend/auth.py
@@ -1,18 +1,19 @@
 from datetime import datetime, timedelta
 from typing import Optional
 
-from config import get_settings
-from database import get_db
 from fastapi import Depends, HTTPException, status
 from fastapi.security import OAuth2PasswordBearer
 from jose import JWTError, jwt
-from models import User as UserModel
 from passlib.context import CryptContext
 from pydantic import BaseModel
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
 
+from config import get_settings
+from database import get_db
+from models import User as UserModel
+
 # Secret key and algorithm for JWT
 SECRET_KEY = get_settings().JWT_SECRET_KEY
 ALGORITHM = get_settings().JWT_ALGORITHM
diff --git a/backend/infra/database.py b/backend/infra/database.py
@@ -1,7 +1,8 @@
-from config import get_settings
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 from sqlalchemy.orm import sessionmaker
 
+from config import get_settings
+
 settings = get_settings()
 
 DATABASE_URL = settings.POSTGRES_DB_URL
diff --git a/backend/infra/redis.py b/backend/infra/redis.py
@@ -1,4 +1,5 @@
 import redis.asyncio as redis
+
 from config import get_settings
 
 settings = get_settings()
diff --git a/backend/infra/vault.py b/backend/infra/vault.py
@@ -1,6 +1,7 @@
 import base64
 
 import httpx
+
 from config import get_settings
 
 
diff --git a/backend/main.py b/backend/main.py
@@ -1,6 +1,7 @@
 import uvicorn
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+
 from middleware.rate_limit import RateLimitMiddleware
 from routes import router
 
diff --git a/backend/middleware/rate_limit.py b/backend/middleware/rate_limit.py
@@ -1,10 +1,11 @@
 import time
 
 from fastapi import Request, status
-from infra.redis import redis_client
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import Response
 
+from infra.redis import redis_client
+
 RATE_LIMIT = 20  # max requests
 RATE_PERIOD = 60  # per 60 seconds
 
diff --git a/backend/routes.py b/backend/routes.py
@@ -1,6 +1,16 @@
 import time
 from datetime import timedelta
 
+from fastapi import APIRouter, Depends, HTTPException, Request, status
+from fastapi.security import OAuth2PasswordRequestForm
+from google.adk import Runner
+from google.adk.sessions import InMemorySessionService
+from google.genai import types
+from pydantic import BaseModel, EmailStr
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.future import select
+
 from agents.query_agent.agent import root_agent
 from auth import (
     ACCESS_TOKEN_EXPIRE_MINUTES,
@@ -11,22 +21,13 @@
     get_password_hash,
 )
 from database import get_db
-from fastapi import APIRouter, Depends, HTTPException, Request, status
-from fastapi.security import OAuth2PasswordRequestForm
-from google.adk import Runner
-from google.adk.sessions import InMemorySessionService
-from google.genai import types
 from middleware.account_lockout import (
     is_account_locked,
     record_failed_login,
     reset_failed_logins,
 )
 from models import User as UserModel
-from pydantic import BaseModel, EmailStr
 from schemas.query_schemas import QueryRequest
-from sqlalchemy.exc import IntegrityError
-from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.future import select
 from utils.helpers import parse_json_markdown
 
 router = APIRouter()
@@ -74,7 +75,7 @@ async def login(
 async def run_query(request: QueryRequest, user=Depends(get_current_active_user)):
     # Setup ADK session
     session_service = InMemorySessionService()
-    session = session_service.create_session(
+    session = await session_service.create_session(
         app_name="QueryApp",
         user_id=user.username,
         session_id=f"{user.username}-session",
diff --git a/backend/scripts/add_sample_queries_qdrant.py b/backend/scripts/add_sample_queries_qdrant.py
@@ -2,11 +2,12 @@
 import json
 import uuid
 
-from config import get_settings
 from litellm import aembedding
 from qdrant_client import AsyncQdrantClient
 from qdrant_client.models import Distance, VectorParams
 
+from config import get_settings
+
 
 async def main(sample_queries_file: str, qdrant_client: AsyncQdrantClient):
     with open(sample_queries_file, "r") as f:
@@ -31,6 +32,11 @@ async def main(sample_queries_file: str, qdrant_client: AsyncQdrantClient):
                 distance=Distance.COSINE,
             ),
         )
+        await qdrant_client.create_payload_index(
+            collection_name=get_settings().QDRANT_COLLECTION_NAME,
+            field_name="type",
+            field_schema="keyword",
+        )
 
     points = []
     for i, embed in enumerate(embeddings_response.data):
diff --git a/backend/scripts/add_table_schemas.py b/backend/scripts/add_table_schemas.py
@@ -2,11 +2,12 @@
 import json
 import uuid
 
-from config import get_settings
 from litellm import aembedding
 from qdrant_client import AsyncQdrantClient
 from qdrant_client.models import Distance, VectorParams
 
+from config import get_settings
+
 
 async def main(tables_file: str, qdrant_client: AsyncQdrantClient):
     with open(tables_file, "r") as f:
@@ -31,6 +32,11 @@ async def main(tables_file: str, qdrant_client: AsyncQdrantClient):
                 distance=Distance.COSINE,
             ),
         )
+        await qdrant_client.create_payload_index(
+            collection_name=get_settings().QDRANT_COLLECTION_NAME,
+            field_name="type",
+            field_schema="keyword",
+        )
 
     points = []
     for i, embed in enumerate(embeddings_response.data):
diff --git a/backend/scripts/clear_vector_db.py b/backend/scripts/clear_vector_db.py
@@ -1,8 +1,9 @@
 import asyncio
 
-from config import get_settings
 from qdrant_client import AsyncQdrantClient
 
+from config import get_settings
+
 
 async def main():
     client = AsyncQdrantClient(
diff --git a/backend/scripts/data/sample_queries.json b/backend/scripts/data/sample_queries.json
diff --git a/backend/scripts/data/tables.json b/backend/scripts/data/tables.json
diff --git a/backend/scripts/generate_sample_queries.py b/backend/scripts/generate_sample_queries.py
diff --git a/backend/scripts/generate_tables_json.py b/backend/scripts/generate_tables_json.py
diff --git a/backend/utils/helpers.py b/backend/utils/helpers.py
