Enhance RelationExtractor with core fixes and verbose logs

- Fix excessive entities being passed to LLM in RelationExtractor
- Add comprehensive 'Heartbeat' verbose logs to methods.py and providers.py
- Ensure robust API key handling and explicit error reporting
This commit is contained in:
KaifAhmad1
2026-01-15 19:00:42 +05:30
parent fd6e058275
commit fb25644fa7
2 changed files with 35 additions and 0 deletions
+11
View File
@@ -1641,6 +1641,11 @@ If a relation doesn't fit any of the preferred types, use the most appropriate t
Extract meaningful relationships between entities. Use appropriate relation types that accurately describe how entities are connected.
Common relation types include: related_to, part_of, located_in, created_by, uses, depends_on, interacts_with, and similar variations."""
verbose_mode = kwargs.get("verbose", False)
if verbose_mode:
import sys
print(f" [methods.extract_relations_llm] Constructing prompt for {len(prompt_entities)} entities...", flush=True, file=sys.stdout)
if not SCHEMAS_AVAILABLE:
raise ImportError("Pydantic schemas not available. Install pydantic/instructor to use LLM extraction.")
@@ -1669,7 +1674,13 @@ Entities found in text: {entities_str}"""
try:
# Use typed generation with Pydantic schema
# Pass kwargs to allow max_tokens and other parameters to be used
if verbose_mode:
import sys
print(f" [methods.extract_relations_llm] Calling llm.generate_typed ({provider}/{model})...", flush=True, file=sys.stdout)
result_obj = llm.generate_typed(prompt, schema=RelationsResponse, **kwargs)
if verbose_mode:
import sys
print(f" [methods.extract_relations_llm] Received response from {provider}.", flush=True, file=sys.stdout)
# Convert back to internal Relation format
relations = []
+24
View File
@@ -318,6 +318,11 @@ class BaseProvider:
"temperature": kwargs.get("temperature", 0.1), # Low temp for structured
}
verbose_mode = kwargs.get("verbose", False)
if verbose_mode:
import sys
print(f" [BaseProvider.generate_typed] Using instructor via {provider_name}. Client: {type(client)}", flush=True, file=sys.stdout)
# Pass through other common parameters
for param in ["max_tokens", "max_completion_tokens", "top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "logit_bias", "user", "top_k"]:
if param in kwargs:
@@ -328,6 +333,9 @@ class BaseProvider:
create_kwargs["response_format"] = {"type": "json_object"}
response = client.chat.completions.create(**create_kwargs)
if verbose_mode:
import sys
print(f" [BaseProvider.generate_typed] Typed response received via instructor ({provider_name}).", flush=True, file=sys.stdout)
return response
except Exception as e:
self.logger.warning(f"Instructor generation failed ({e}), falling back to manual repair loop.")
@@ -706,7 +714,15 @@ class GroqProvider(BaseProvider):
if param in kwargs:
create_kwargs[param] = kwargs[param]
verbose_mode = kwargs.get("verbose", False)
if verbose_mode:
import sys
print(f" [GroqProvider.generate] Sending request to Groq API (model: {create_kwargs['model']})...", flush=True, file=sys.stdout)
response = self.client.chat.completions.create(**create_kwargs)
if verbose_mode:
import sys
print(f" [GroqProvider.generate] Response received from Groq.", flush=True, file=sys.stdout)
return response.choices[0].message.content
def generate_structured(self, prompt: str, **kwargs) -> dict:
@@ -737,7 +753,15 @@ class GroqProvider(BaseProvider):
if param in kwargs:
create_kwargs[param] = kwargs[param]
verbose_mode = kwargs.get("verbose", False)
if verbose_mode:
import sys
print(f" [GroqProvider.generate_structured] Sending structured request to Groq API (model: {create_kwargs['model']})...", flush=True, file=sys.stdout)
response = self.client.chat.completions.create(**create_kwargs)
if verbose_mode:
import sys
print(f" [GroqProvider.generate_structured] Structured response received from Groq.", flush=True, file=sys.stdout)
try:
return self._parse_json(response.choices[0].message.content)
except Exception as e: