Asish Karthikeya Gogineni commited on
Commit
6d5c110
·
1 Parent(s): cddcaaf

feat: Add local embeddings to bypass API rate limits

Browse files
Files changed (2) hide show
  1. app.py +6 -5
  2. code_chatbot/indexer.py +13 -3
app.py CHANGED
@@ -380,17 +380,18 @@ with st.sidebar:
380
  os.environ["QDRANT_API_KEY"] = qdrant_key
381
 
382
  # For Groq, we need an embedding provider
383
- embedding_provider = provider
 
384
  embedding_api_key = api_key
385
 
386
  if provider == "groq":
387
- st.info(f"ℹ️ {provider.capitalize()} is used for Chat. For indexing, please select 'gemini' for embeddings.")
388
- embedding_provider = "gemini" # Force gemini if groq is used, as openai is removed
389
 
390
- # Check Embedding Key for Gemini
391
  emb_env_key = os.getenv("GOOGLE_API_KEY")
392
  if not emb_env_key and provider != "gemini":
393
- embedding_api_key = st.text_input("Google API Key (for Embeddings)", type="password")
394
  else:
395
  embedding_api_key = emb_env_key
396
 
 
380
  os.environ["QDRANT_API_KEY"] = qdrant_key
381
 
382
  # For Groq, we need an embedding provider
383
+ # Use LOCAL embeddings by default - NO RATE LIMITS!
384
+ embedding_provider = "local" # Use local HuggingFace embeddings
385
  embedding_api_key = api_key
386
 
387
  if provider == "groq":
388
+ st.info(f"ℹ️ {provider.capitalize()} is used for Chat. Using LOCAL embeddings (no rate limits!).")
389
+ embedding_provider = "local" # Use local embeddings for Groq too
390
 
391
+ # Check Embedding Key for Gemini (not needed for local)
392
  emb_env_key = os.getenv("GOOGLE_API_KEY")
393
  if not emb_env_key and provider != "gemini":
394
+ embedding_api_key = emb_env_key # Optional now
395
  else:
396
  embedding_api_key = emb_env_key
397
 
code_chatbot/indexer.py CHANGED
@@ -62,11 +62,20 @@ class Indexer:
62
  )
63
  logger.info("Path obfuscation enabled")
64
 
65
- # Setup Embeddings (only Gemini supported)
66
  if embedding_function:
67
  self.embedding_function = embedding_function
68
  else:
69
- if provider == "gemini":
 
 
 
 
 
 
 
 
 
70
  api_key = api_key or os.getenv("GOOGLE_API_KEY")
71
  if not api_key:
72
  raise ValueError("Google API Key is required for Gemini Embeddings")
@@ -74,8 +83,9 @@ class Indexer:
74
  model="models/gemini-embedding-001",
75
  google_api_key=api_key
76
  )
 
77
  else:
78
- raise ValueError(f"Unsupported embedding provider: {provider}. Only 'gemini' is supported.")
79
 
80
  def clear_collection(self, collection_name: str = "codebase"):
81
  """
 
62
  )
63
  logger.info("Path obfuscation enabled")
64
 
65
+ # Setup Embeddings - supports Gemini (API) and local HuggingFace
66
  if embedding_function:
67
  self.embedding_function = embedding_function
68
  else:
69
+ if provider == "local" or provider == "huggingface":
70
+ # Use local embeddings - NO RATE LIMITS!
71
+ from langchain_huggingface import HuggingFaceEmbeddings
72
+ self.embedding_function = HuggingFaceEmbeddings(
73
+ model_name="all-MiniLM-L6-v2", # Fast & good quality
74
+ model_kwargs={'device': 'cpu'},
75
+ encode_kwargs={'normalize_embeddings': True}
76
+ )
77
+ logger.info("Using LOCAL embeddings (no rate limits)")
78
+ elif provider == "gemini":
79
  api_key = api_key or os.getenv("GOOGLE_API_KEY")
80
  if not api_key:
81
  raise ValueError("Google API Key is required for Gemini Embeddings")
 
83
  model="models/gemini-embedding-001",
84
  google_api_key=api_key
85
  )
86
+ logger.info("Using Gemini embeddings (API rate limits apply)")
87
  else:
88
+ raise ValueError(f"Unsupported embedding provider: {provider}. Use 'local', 'huggingface', or 'gemini'.")
89
 
90
  def clear_collection(self, collection_name: str = "codebase"):
91
  """